diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index a86b891..bf8fd93 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,23 +1,26 @@ +minimum_pre_commit_version: "4.6.0" + repos: - repo: https://github.com/pre-commit/pre-commit-hooks - rev: v4.5.0 + rev: v6.0.0 hooks: - id: check-merge-conflict - id: check-yaml - id: end-of-file-fixer - id: trailing-whitespace - - id: end-of-file-fixer - - id: check-yaml - repo: https://github.com/PyCQA/isort - rev: 5.12.0 + rev: 8.0.1 hooks: - id: isort name: isort - repo: https://github.com/pre-commit/mirrors-mypy - rev: v1.7.1 + rev: v1.20.2 hooks: - id: mypy + additional_dependencies: [types-requests] + files: ^fedgraph/ + args: [--config-file=pyproject.toml] - repo: https://github.com/psf/black - rev: 23.11.0 + rev: 26.5.1 hooks: - id: black diff --git a/benchmark/benchmark_NC_FedGraph.py b/benchmark/benchmark_NC_FedGraph.py index 256aa07..c264d55 100644 --- a/benchmark/benchmark_NC_FedGraph.py +++ b/benchmark/benchmark_NC_FedGraph.py @@ -69,8 +69,8 @@ config = attridict(config) run_fedgraph(config) -config.iid_beta=100 +config.iid_beta = 100 run_fedgraph(config) -config.iid_beta=10 -run_fedgraph(config) \ No newline at end of file +config.iid_beta = 10 +run_fedgraph(config) diff --git a/benchmark/benchmark_NC_FederatedScope.py b/benchmark/benchmark_NC_FederatedScope.py index 236e67c..eb191f5 100644 --- a/benchmark/benchmark_NC_FederatedScope.py +++ b/benchmark/benchmark_NC_FederatedScope.py @@ -67,7 +67,7 @@ def load_data(config, client_cfgs=None): ds = Planetoid(root="data/", name=PLANETOID_NAMES[name]) full = ds[0] num_classes = int(full.y.max().item()) + 1 - + # 与data_process.py完全一致：在全部节点上做Dirichlet分割 split_idxs = label_dirichlet_partition( full.y, # 使用全部节点，不只是训练节点 @@ -77,16 +77,16 @@ def load_data(config, client_cfgs=None): config.iid_beta, config.distribution_type, ) - + parts = [] for idxs in split_idxs: client_nodes = torch.tensor(idxs) - + # 为每个客户端创建mask，但保持原有的train/val/test划分逻辑 train_mask = torch.zeros(full.num_nodes, dtype=torch.bool) val_mask = torch.zeros(full.num_nodes, dtype=torch.bool) test_mask = torch.zeros(full.num_nodes, dtype=torch.bool) - + # 在客户端节点中，保持原有数据集的train/val/test划分 for node in client_nodes: if full.train_mask[node]: @@ -95,20 +95,21 @@ def load_data(config, client_cfgs=None): val_mask[node] = True elif full.test_mask[node]: test_mask[node] = True - + parts.append( Data( x=full.x, edge_index=full.edge_index, y=full.y, train_mask=train_mask, # 保持原有train划分 - val_mask=val_mask, # 保持原有val划分 - test_mask=test_mask, # 保持原有test划分 + val_mask=val_mask, # 保持原有val划分 + test_mask=test_mask, # 保持原有test划分 ) ) - + data_dict = { - i + 1: { + i + + 1: { "data": parts[i], "train": [parts[i]], "val": [parts[i]], @@ -139,7 +140,6 @@ def build(cfg_model, input_shape): register_model(mkey, builder) - def run_fedscope_experiment(ds, beta): cfg = global_cfg.clone() cfg.defrost() @@ -150,7 +150,7 @@ def run_fedscope_experiment(ds, beta): cfg.federate.client_num = CLIENT_NUM cfg.federate.total_round_num = TOTAL_ROUNDS cfg.federate.make_global_eval = False - cfg.federate.process_num = CLIENT_NUM + cfg.federate.process_num = CLIENT_NUM cfg.federate.num_cpus_per_trainer = CPUS_PER_TRAINER cfg.data.root = "data/" cfg.data.type = ds @@ -180,13 +180,17 @@ def run_fedscope_experiment(ds, beta): res = runner.run() dur = time.time() - t0 mem = peak_memory_mb() - + # 获取FederatedScope结果 - + # 从FederatedScope的结果中获取准确率 # 使用加权平均以与FedGraph保持一致 - acc = res.get("client_summarized_weighted_avg", {}).get("test_acc", 0.0) if res else 0.0 - + acc = ( + res.get("client_summarized_weighted_avg", {}).get("test_acc", 0.0) + if res + else 0.0 + ) + acc_pct = acc * 100 if acc <= 1.0 else acc model = runner.server.model if runner.server else None if model is not None: @@ -210,8 +214,6 @@ def run_fedscope_experiment(ds, beta): } - - def main(): parser = argparse.ArgumentParser() parser.add_argument("--use_cluster", action="store_true") diff --git a/benchmark/benchmark_NC_old.py b/benchmark/benchmark_NC_old.py index f9fcb01..352de2c 100644 --- a/benchmark/benchmark_NC_old.py +++ b/benchmark/benchmark_NC_old.py @@ -18,9 +18,9 @@ # Datasets to benchmark datasets = [ "cora", - #"citeseer", - #"pubmed", - #"ogbn-arxiv", + # "citeseer", + # "pubmed", + # "ogbn-arxiv", ] # You can add more: ["cora", "citeseer", "ogbn-arxiv", "ogbn-products"] # datasets = ["ogbn-papers100M"] # Number of trainers to test diff --git a/benchmark/exp/config.yaml b/benchmark/exp/config.yaml index 74d0faa..82d11c4 100644 --- a/benchmark/exp/config.yaml +++ b/benchmark/exp/config.yaml @@ -372,4 +372,3 @@ vertical: use: false wandb: use: false - diff --git a/benchmark/figure/GC_comm_costs/extract_GC_log.py b/benchmark/figure/GC_comm_costs/extract_GC_log.py index 86a2407..d09272b 100644 --- a/benchmark/figure/GC_comm_costs/extract_GC_log.py +++ b/benchmark/figure/GC_comm_costs/extract_GC_log.py @@ -141,18 +141,18 @@ def extract_metrics(exp_text, algorithm, dataset, trainers): "Trainers": trainers, "Accuracy": accuracy, "Train_Time_ms": train_time, - "Theoretical_Pretrain_MB": float(theoretical_pretrain[-1]) - if theoretical_pretrain - else 0, - "Theoretical_Train_MB": float(theoretical_train[-1]) - if theoretical_train - else 0, - "Actual_Pretrain_MB": float(actual_pretrain_match.group(1)) - if actual_pretrain_match - else None, - "Actual_Train_MB": float(actual_train_match.group(1)) - if actual_train_match - else None, + "Theoretical_Pretrain_MB": ( + float(theoretical_pretrain[-1]) if theoretical_pretrain else 0 + ), + "Theoretical_Train_MB": ( + float(theoretical_train[-1]) if theoretical_train else 0 + ), + "Actual_Pretrain_MB": ( + float(actual_pretrain_match.group(1)) if actual_pretrain_match else None + ), + "Actual_Train_MB": ( + float(actual_train_match.group(1)) if actual_train_match else None + ), } # Calculate totals @@ -185,15 +185,19 @@ def generate_accuracy_comparison(df, output_file="gc_accuracy_comparison.pdf"): plt.figure(figsize=(14, 8)) datasets = sorted( comparison_data["Dataset"].unique(), - key=lambda x: ["IMDB-BINARY", "IMDB-MULTI", "MUTAG", "BZR", "COX2"].index(x) - if x in ["IMDB-BINARY", "IMDB-MULTI", "MUTAG", "BZR", "COX2"] - else 999, + key=lambda x: ( + ["IMDB-BINARY", "IMDB-MULTI", "MUTAG", "BZR", "COX2"].index(x) + if x in ["IMDB-BINARY", "IMDB-MULTI", "MUTAG", "BZR", "COX2"] + else 999 + ), ) algorithms = sorted( comparison_data["Algorithm"].unique(), - key=lambda x: ["FedAvg", "GCFL", "GCFL+", "GCFL+dWs"].index(x) - if x in ["FedAvg", "GCFL", "GCFL+", "GCFL+dWs"] - else 999, + key=lambda x: ( + ["FedAvg", "GCFL", "GCFL+", "GCFL+dWs"].index(x) + if x in ["FedAvg", "GCFL", "GCFL+", "GCFL+dWs"] + else 999 + ), ) x_positions = np.arange(len(datasets)) width = 0.8 / len(algorithms) @@ -249,15 +253,19 @@ def generate_train_time_comparison(df, output_file="gc_train_time_comparison.pdf plt.figure(figsize=(14, 8)) datasets = sorted( comparison_data["Dataset"].unique(), - key=lambda x: ["IMDB-BINARY", "IMDB-MULTI", "MUTAG", "BZR", "COX2"].index(x) - if x in ["IMDB-BINARY", "IMDB-MULTI", "MUTAG", "BZR", "COX2"] - else 999, + key=lambda x: ( + ["IMDB-BINARY", "IMDB-MULTI", "MUTAG", "BZR", "COX2"].index(x) + if x in ["IMDB-BINARY", "IMDB-MULTI", "MUTAG", "BZR", "COX2"] + else 999 + ), ) algorithms = sorted( comparison_data["Algorithm"].unique(), - key=lambda x: ["FedAvg", "GCFL", "GCFL+", "GCFL+dWs"].index(x) - if x in ["FedAvg", "GCFL", "GCFL+", "GCFL+dWs"] - else 999, + key=lambda x: ( + ["FedAvg", "GCFL", "GCFL+", "GCFL+dWs"].index(x) + if x in ["FedAvg", "GCFL", "GCFL+", "GCFL+dWs"] + else 999 + ), ) x_positions = np.arange(len(datasets)) width = 0.8 / len(algorithms) @@ -330,16 +338,20 @@ def generate_comm_cost_comparison(df, output_file="gc_comm_cost_comparison.pdf") # Datasets and algorithms datasets = sorted( comparison_data["Dataset"].unique(), - key=lambda x: ["IMDB-BINARY", "IMDB-MULTI", "MUTAG", "BZR", "COX2"].index(x) - if x in ["IMDB-BINARY", "IMDB-MULTI", "MUTAG", "BZR", "COX2"] - else 999, + key=lambda x: ( + ["IMDB-BINARY", "IMDB-MULTI", "MUTAG", "BZR", "COX2"].index(x) + if x in ["IMDB-BINARY", "IMDB-MULTI", "MUTAG", "BZR", "COX2"] + else 999 + ), ) algorithms = sorted( comparison_data["Algorithm"].unique(), - key=lambda x: ["FedAvg", "GCFL", "GCFL+", "GCFL+dWs"].index(x) - if x in ["FedAvg", "GCFL", "GCFL+", "GCFL+dWs"] - else 999, + key=lambda x: ( + ["FedAvg", "GCFL", "GCFL+", "GCFL+dWs"].index(x) + if x in ["FedAvg", "GCFL", "GCFL+", "GCFL+dWs"] + else 999 + ), ) # X-axis setup diff --git a/benchmark/figure/GC_comm_costs_old/extract_GC_log_old.py b/benchmark/figure/GC_comm_costs_old/extract_GC_log_old.py index 2bc20f1..f2e073c 100644 --- a/benchmark/figure/GC_comm_costs_old/extract_GC_log_old.py +++ b/benchmark/figure/GC_comm_costs_old/extract_GC_log_old.py @@ -141,18 +141,18 @@ def extract_metrics(exp_text, algorithm, dataset, trainers): "Trainers": trainers, "Accuracy": accuracy, "Train_Time_ms": train_time, - "Theoretical_Pretrain_MB": float(theoretical_pretrain[-1]) - if theoretical_pretrain - else 0, - "Theoretical_Train_MB": float(theoretical_train[-1]) - if theoretical_train - else 0, - "Actual_Pretrain_MB": float(actual_pretrain_match.group(1)) - if actual_pretrain_match - else None, - "Actual_Train_MB": float(actual_train_match.group(1)) - if actual_train_match - else None, + "Theoretical_Pretrain_MB": ( + float(theoretical_pretrain[-1]) if theoretical_pretrain else 0 + ), + "Theoretical_Train_MB": ( + float(theoretical_train[-1]) if theoretical_train else 0 + ), + "Actual_Pretrain_MB": ( + float(actual_pretrain_match.group(1)) if actual_pretrain_match else None + ), + "Actual_Train_MB": ( + float(actual_train_match.group(1)) if actual_train_match else None + ), } # Calculate totals @@ -195,16 +195,20 @@ def generate_accuracy_comparison(df, output_file="gc_accuracy_comparison.pdf"): # Get unique datasets and algorithms in desired order datasets = sorted( comparison_data["Dataset"].unique(), - key=lambda x: ["IMDB-BINARY", "IMDB-MULTI", "MUTAG", "BZR", "COX2"].index(x) - if x in ["IMDB-BINARY", "IMDB-MULTI", "MUTAG", "BZR", "COX2"] - else 999, + key=lambda x: ( + ["IMDB-BINARY", "IMDB-MULTI", "MUTAG", "BZR", "COX2"].index(x) + if x in ["IMDB-BINARY", "IMDB-MULTI", "MUTAG", "BZR", "COX2"] + else 999 + ), ) algorithms = sorted( comparison_data["Algorithm"].unique(), - key=lambda x: ["FedAvg", "GCFL", "GCFL+", "GCFL+dWs"].index(x) - if x in ["FedAvg", "GCFL", "GCFL+", "GCFL+dWs"] - else 999, + key=lambda x: ( + ["FedAvg", "GCFL", "GCFL+", "GCFL+dWs"].index(x) + if x in ["FedAvg", "GCFL", "GCFL+", "GCFL+dWs"] + else 999 + ), ) # Set x positions @@ -289,16 +293,20 @@ def generate_train_time_comparison(df, output_file="gc_train_time_comparison.pdf # Get unique datasets and algorithms in desired order datasets = sorted( comparison_data["Dataset"].unique(), - key=lambda x: ["IMDB-BINARY", "IMDB-MULTI", "MUTAG", "BZR", "COX2"].index(x) - if x in ["IMDB-BINARY", "IMDB-MULTI", "MUTAG", "BZR", "COX2"] - else 999, + key=lambda x: ( + ["IMDB-BINARY", "IMDB-MULTI", "MUTAG", "BZR", "COX2"].index(x) + if x in ["IMDB-BINARY", "IMDB-MULTI", "MUTAG", "BZR", "COX2"] + else 999 + ), ) algorithms = sorted( comparison_data["Algorithm"].unique(), - key=lambda x: ["FedAvg", "GCFL", "GCFL+", "GCFL+dWs"].index(x) - if x in ["FedAvg", "GCFL", "GCFL+", "GCFL+dWs"] - else 999, + key=lambda x: ( + ["FedAvg", "GCFL", "GCFL+", "GCFL+dWs"].index(x) + if x in ["FedAvg", "GCFL", "GCFL+", "GCFL+dWs"] + else 999 + ), ) # Set x positions @@ -390,16 +398,20 @@ def generate_comm_cost_comparison(df, output_file="gc_comm_cost_comparison.pdf") # Datasets and algorithms datasets = sorted( comparison_data["Dataset"].unique(), - key=lambda x: ["IMDB-BINARY", "IMDB-MULTI", "MUTAG", "BZR", "COX2"].index(x) - if x in ["IMDB-BINARY", "IMDB-MULTI", "MUTAG", "BZR", "COX2"] - else 999, + key=lambda x: ( + ["IMDB-BINARY", "IMDB-MULTI", "MUTAG", "BZR", "COX2"].index(x) + if x in ["IMDB-BINARY", "IMDB-MULTI", "MUTAG", "BZR", "COX2"] + else 999 + ), ) algorithms = sorted( comparison_data["Algorithm"].unique(), - key=lambda x: ["FedAvg", "GCFL", "GCFL+", "GCFL+dWs"].index(x) - if x in ["FedAvg", "GCFL", "GCFL+", "GCFL+dWs"] - else 999, + key=lambda x: ( + ["FedAvg", "GCFL", "GCFL+", "GCFL+dWs"].index(x) + if x in ["FedAvg", "GCFL", "GCFL+", "GCFL+dWs"] + else 999 + ), ) # X-axis setup diff --git a/benchmark/figure/LP_comm_costs/extract_LP_log.py b/benchmark/figure/LP_comm_costs/extract_LP_log.py index 9b94eb9..6bde527 100644 --- a/benchmark/figure/LP_comm_costs/extract_LP_log.py +++ b/benchmark/figure/LP_comm_costs/extract_LP_log.py @@ -122,18 +122,22 @@ def extract_lp_data(logfile): "Dataset": countries, "AUC": auc, "TrainTime": train_time_ms, - "Theoretical_Pretrain_MB": float(theoretical_pretrain_match.group(1)) - if theoretical_pretrain_match - else 0, - "Theoretical_Train_MB": float(theoretical_train_match.group(1)) - if theoretical_train_match - else 0, - "Actual_Pretrain_MB": float(actual_pretrain_match.group(1)) - if actual_pretrain_match - else 0, - "Actual_Train_MB": float(actual_train_match.group(1)) - if actual_train_match - else 0, + "Theoretical_Pretrain_MB": ( + float(theoretical_pretrain_match.group(1)) + if theoretical_pretrain_match + else 0 + ), + "Theoretical_Train_MB": ( + float(theoretical_train_match.group(1)) + if theoretical_train_match + else 0 + ), + "Actual_Pretrain_MB": ( + float(actual_pretrain_match.group(1)) if actual_pretrain_match else 0 + ), + "Actual_Train_MB": ( + float(actual_train_match.group(1)) if actual_train_match else 0 + ), "Hit_Rate": hit_rate, } diff --git a/benchmark/figure/NC_comm_costs/extract_NC_100M_log.py b/benchmark/figure/NC_comm_costs/extract_NC_100M_log.py index 74fd49a..8e57958 100644 --- a/benchmark/figure/NC_comm_costs/extract_NC_100M_log.py +++ b/benchmark/figure/NC_comm_costs/extract_NC_100M_log.py @@ -7,6 +7,7 @@ Dataset: <#nodes> nodes, <#edges> edges [ β=] Round → Test Acc: % | Computation Time: s | Memory: MB | Comm Cost: MB """ + import logging import warnings diff --git a/benchmark/figure/NC_comm_costs/extract_NC_log.py b/benchmark/figure/NC_comm_costs/extract_NC_log.py index 69f58e3..d750e09 100644 --- a/benchmark/figure/NC_comm_costs/extract_NC_log.py +++ b/benchmark/figure/NC_comm_costs/extract_NC_log.py @@ -92,18 +92,18 @@ def extract_metrics(exp_text, algorithm, dataset, trainers, iid_beta): "IID_Beta": iid_beta, "Accuracy": accuracy, "Train_Time_ms": train_time, - "Theoretical_Pretrain_MB": float(theoretical_pretrain[-1]) - if theoretical_pretrain - else 0, - "Theoretical_Train_MB": float(theoretical_train[-1]) - if theoretical_train - else 0, - "Actual_Pretrain_MB": float(actual_pretrain_match.group(1)) - if actual_pretrain_match - else None, - "Actual_Train_MB": float(actual_train_match.group(1)) - if actual_train_match - else None, + "Theoretical_Pretrain_MB": ( + float(theoretical_pretrain[-1]) if theoretical_pretrain else 0 + ), + "Theoretical_Train_MB": ( + float(theoretical_train[-1]) if theoretical_train else 0 + ), + "Actual_Pretrain_MB": ( + float(actual_pretrain_match.group(1)) if actual_pretrain_match else None + ), + "Actual_Train_MB": ( + float(actual_train_match.group(1)) if actual_train_match else None + ), } result["Theoretical_Total_MB"] = ( result["Theoretical_Pretrain_MB"] + result["Theoretical_Train_MB"] diff --git a/benchmark/figure/NC_comm_costs_old/extract_NC_log_old.py b/benchmark/figure/NC_comm_costs_old/extract_NC_log_old.py index 76f86ef..ddfbb49 100644 --- a/benchmark/figure/NC_comm_costs_old/extract_NC_log_old.py +++ b/benchmark/figure/NC_comm_costs_old/extract_NC_log_old.py @@ -94,18 +94,18 @@ def extract_metrics(exp_text, algorithm, dataset, trainers, iid_beta): "Accuracy": accuracy, "Train_Time_ms": train_time_ms, "Train_Time_s": train_time_s, - "Theoretical_Pretrain_MB": float(theoretical_pretrain[-1]) - if theoretical_pretrain - else 0, - "Theoretical_Train_MB": float(theoretical_train[-1]) - if theoretical_train - else 0, - "Actual_Pretrain_MB": float(actual_pretrain_match.group(1)) - if actual_pretrain_match - else None, - "Actual_Train_MB": float(actual_train_match.group(1)) - if actual_train_match - else None, + "Theoretical_Pretrain_MB": ( + float(theoretical_pretrain[-1]) if theoretical_pretrain else 0 + ), + "Theoretical_Train_MB": ( + float(theoretical_train[-1]) if theoretical_train else 0 + ), + "Actual_Pretrain_MB": ( + float(actual_pretrain_match.group(1)) if actual_pretrain_match else None + ), + "Actual_Train_MB": ( + float(actual_train_match.group(1)) if actual_train_match else None + ), } result["Theoretical_Total_MB"] = ( result["Theoretical_Pretrain_MB"] + result["Theoretical_Train_MB"] diff --git a/docs/dev_script/processing_script_GC.py b/docs/dev_script/processing_script_GC.py index e20df8c..aaee1ca 100644 --- a/docs/dev_script/processing_script_GC.py +++ b/docs/dev_script/processing_script_GC.py @@ -78,9 +78,9 @@ def process_log(log_content): train_network_match = re.search(r"Log ([^,]+) network: (\d+\.\d+)", line) if train_network_match and train_mode: - current_experiment[ - f"Train Network {(train_network_match.group(1))}" - ] = float(train_network_match.group(2)) + current_experiment[f"Train Network {(train_network_match.group(1))}"] = ( + float(train_network_match.group(2)) + ) average_accuracy_match = re.search(r"Average test accuracy: (\d+\.\d+)", line) if average_accuracy_match: current_experiment["Average Test Accuracy"] = float( diff --git a/docs/dev_script/processing_script_LP.py b/docs/dev_script/processing_script_LP.py index 133cdc6..3ebfc2d 100644 --- a/docs/dev_script/processing_script_LP.py +++ b/docs/dev_script/processing_script_LP.py @@ -81,9 +81,9 @@ def process_log(log_content): train_network_match = re.search(r"Log ([^,]+) network: (\d+\.\d+)", line) if train_network_match and train_mode: - current_experiment[ - f"Train Network {(train_network_match.group(1))}" - ] = float(train_network_match.group(2)) + current_experiment[f"Train Network {(train_network_match.group(1))}"] = ( + float(train_network_match.group(2)) + ) average_accuracy_match = re.search( r"Predict Day 20 average auc score: (\d+\.\d+) hit rate: (\d+\.\d+)", line ) diff --git a/docs/dev_script/processing_script_NC.py b/docs/dev_script/processing_script_NC.py index dff3ed7..72c973e 100644 --- a/docs/dev_script/processing_script_NC.py +++ b/docs/dev_script/processing_script_NC.py @@ -81,9 +81,9 @@ def process_log(log_content): train_network_match = re.search(r"Log ([^,]+) network: (\d+\.\d+)", line) if train_network_match and train_mode: - current_experiment[ - f"Train Network {(train_network_match.group(1))}" - ] = float(train_network_match.group(2)) + current_experiment[f"Train Network {(train_network_match.group(1))}"] = ( + float(train_network_match.group(2)) + ) average_accuracy_match = re.search(r"Average test accuracy: (\d+\.\d+)", line) if average_accuracy_match: current_experiment["Average Test Accuracy"] = float( diff --git a/docs/index.rst b/docs/index.rst index ac7ce86..443048a 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -79,3 +79,4 @@ Whether you are a federated learning researcher or a first-time user of federate cite reference + pre_commit diff --git a/docs/pre_commit.rst b/docs/pre_commit.rst new file mode 100644 index 0000000..9947eaf --- /dev/null +++ b/docs/pre_commit.rst @@ -0,0 +1,94 @@ +Pre-commit Workflow +=================== + +FedGraph uses `pre-commit `__ to run lightweight +repository checks before a commit is created. These checks are a code-quality +gate; they do not run the test suite in ``tests/``. Run tests separately with +``pytest`` when a change needs behavioral verification. + +Setup and Trigger Points +------------------------ + +Install the development tools and the Git hook once in a checkout:: + + python -m pip install -e '.[dev]' + python -m pre_commit install --install-hooks + +After that, ``git commit`` invokes ``.git/hooks/pre-commit``. The hook reads +``.pre-commit-config.yaml`` and runs its checks on staged files. To run the +same checks over the whole checkout, for example after updating hooks, use:: + + python -m pre_commit run --all-files + +The hook environments are cached independently by pre-commit, so the first run +downloads and creates them; later runs reuse the cache. Continuous integration +can use the same ``run --all-files`` command to apply the same gate to a clean +checkout. + +Configuration +------------- + +``.pre-commit-config.yaml`` is the entry point. Its +``minimum_pre_commit_version`` requires a runner new enough to interpret this +configuration. Each ``repos`` item pins a hook repository revision, which +makes developers and CI use the same checker versions. + +The configured hooks are: + +* ``check-merge-conflict``: rejects unresolved Git conflict markers such as + ``<<<<<<<`` and ``>>>>>>>``. +* ``check-yaml``: parses YAML files and reports invalid syntax. +* ``end-of-file-fixer``: ensures text files end with exactly one newline. +* ``trailing-whitespace``: removes whitespace at the ends of lines. +* ``isort``: orders and groups Python imports. Its layout follows Black's + formatting profile, configured in ``pyproject.toml``. +* ``mypy``: performs static type checking for ``fedgraph/`` only. It installs + ``types-requests`` in its isolated hook environment because the package uses + ``requests`` types. +* ``black``: applies one consistent Python code format. + +The final three hooks may rewrite files. Review and stage those edits, then run +the commit again. The first four are repository-hygiene checks. + +Related Project Files +--------------------- + +``pyproject.toml`` supplies tool-specific settings read by the hooks: + +* ``[tool.isort] profile = "black"`` prevents isort and Black from disagreeing + about import formatting. +* ``[tool.mypy] python_version = "3.9"`` checks against FedGraph's supported + minimum Python version, while ``ignore_missing_imports = true`` avoids + treating untyped optional third-party packages as project errors. + +``setup.py`` defines the optional ``dev`` dependency group. Installing +``.[dev]`` provides the pre-commit runner, a compatible Mypy release, and +pytest for local development. Pre-commit itself creates isolated environments +from the exact revisions in ``.pre-commit-config.yaml``; ``setup.py`` makes the +runner and normal development tools available in the developer's environment. + +Typing Changes in Core Source +----------------------------- + +Updating Mypy exposed a small set of static-type ambiguities in the core +implementation. The following changes are annotations or narrowly scoped +``# type: ignore`` markers; they do not alter FedGraph's runtime behavior. + +* ``fedgraph/server_class.py`` annotates aggregation statistics and handles the + optional OpenFHE class import. +* ``fedgraph/trainer_class.py`` marks optional constructor metadata, handles + the optional OpenFHE import, and documents model/context attributes that are + initialized later in the trainer lifecycle. +* ``fedgraph/federated_methods.py`` marks optional OpenFHE availability and + Ray's runtime-created ``Trainer.remote`` API, which static analysis cannot + infer from the decorator. +* ``fedgraph/low_rank/server_lowrank.py`` annotates compression statistics. +* ``fedgraph/low_rank/trainer_lowrank.py`` documents model attributes inherited + from the base trainer and initialized after construction. +* ``fedgraph/utils_nc.py`` gives batched node-index partitions their nested + integer-list type. +* ``fedgraph/gnn_models.py`` documents that GIN's optional layers are fully + configured before its training-time ``forward`` path executes. + +Black and isort also reformatted other Python files. Those mechanical changes +are separate from the typing adjustments above. diff --git a/fedgraph/__init__.py b/fedgraph/__init__.py index 32494a2..d1560f6 100644 --- a/fedgraph/__init__.py +++ b/fedgraph/__init__.py @@ -3,12 +3,12 @@ federated_methods, gnn_models, monitor_class, + openfhe_threshold, server_class, train_func, trainer_class, utils_gc, utils_lp, utils_nc, - openfhe_threshold, ) from .version import __version__ diff --git a/fedgraph/data_process.py b/fedgraph/data_process.py index addbf9b..803995d 100644 --- a/fedgraph/data_process.py +++ b/fedgraph/data_process.py @@ -353,7 +353,9 @@ def NC_load_data(dataset_str: str) -> tuple: # large for HE simulation (2.4M nodes). sub_n = int(os.environ.get("FEDGRAPH_SUBSAMPLE_NODES", "0")) if sub_n > 0 and features.shape[0] > sub_n: - print(f"[subsample] Reducing {dataset_str} from {features.shape[0]} to {sub_n} nodes") + print( + f"[subsample] Reducing {dataset_str} from {features.shape[0]} to {sub_n} nodes" + ) keep = torch.arange(sub_n) keep_set = set(keep.tolist()) features = features[keep] @@ -362,7 +364,8 @@ def NC_load_data(dataset_str: str) -> tuple: row, col, val = adj.coo() edge_mask = (row < sub_n) & (col < sub_n) adj = torch_sparse.SparseTensor( - row=row[edge_mask], col=col[edge_mask], + row=row[edge_mask], + col=col[edge_mask], value=val[edge_mask] if val is not None else None, sparse_sizes=(sub_n, sub_n), ) @@ -370,8 +373,10 @@ def NC_load_data(dataset_str: str) -> tuple: idx_train = idx_train[idx_train < sub_n] idx_val = idx_val[idx_val < sub_n] idx_test = idx_test[idx_test < sub_n] - print(f"[subsample] kept {features.shape[0]} nodes, {edge_mask.sum().item()} edges, " - f"{len(idx_train)} train, {len(idx_test)} test") + print( + f"[subsample] kept {features.shape[0]} nodes, {edge_mask.sum().item()} edges, " + f"{len(idx_train)} train, {len(idx_test)} test" + ) elif dataset_str == "reddit": from dgl.data import RedditDataset diff --git a/fedgraph/federated_methods.py b/fedgraph/federated_methods.py index a6ee7c2..64360b5 100644 --- a/fedgraph/federated_methods.py +++ b/fedgraph/federated_methods.py @@ -1,4 +1,3 @@ - import argparse import copy import datetime @@ -39,7 +38,7 @@ try: from fedgraph.openfhe_threshold import OpenFHEThresholdCKKS # noqa: F401 except ImportError: # pragma: no cover - OpenFHEThresholdCKKS = None # type: ignore[assignment] + OpenFHEThresholdCKKS = None # type: ignore[misc,assignment] # Optional differential privacy module. try: @@ -168,9 +167,7 @@ def run_fedgraph(args: attridict) -> None: # the FedAvg-and-no-encryption combination from the previous prototype is # one supported regime but not a requirement. if getattr(args, "use_lowrank", False) and args.fedgraph_task != "NC": - raise ValueError( - "Low-rank compression currently only supported for NC tasks" - ) + raise ValueError("Low-rank compression currently only supported for NC tasks") if args.fedgraph_task == "NC": _validate_nc_num_hops(args) @@ -482,9 +479,7 @@ def get_memory_usage(self): use_lowrank = getattr(args, "use_lowrank", False) and LOWRANK_AVAILABLE ServerClass = Server_LowRank if use_lowrank else Server _feature_dim = feature_shape if args.use_huggingface else features.shape[1] - server = ServerClass( - _feature_dim, args_hidden, class_num, device, trainers, args - ) + server = ServerClass(_feature_dim, args_hidden, class_num, device, trainers, args) # End initialization time tracking server.broadcast_params(-1) @@ -509,7 +504,9 @@ def get_memory_usage(self): ] results = ray.get(encrypted_data) - encrypted_sums = [(r[0], r[1]) for r in results] # (encrypted_sum, shape) + encrypted_sums = [ + (r[0], r[1]) for r in results + ] # (encrypted_sum, shape) encryption_times = [r[2] for r in results] enc_sizes = [len(r[0]) for r in results] # size of encrypted data @@ -538,8 +535,10 @@ def get_memory_usage(self): decryption_times = ray.get(load_feature_refs) elif getattr(args, "he_backend", "tenseal") == "openfhe": print("Starting OpenFHE threshold encrypted feature aggregation...") + import os + import tempfile + import openfhe - import os, tempfile # Shared directory for OpenFHE serialized objects he_dir = tempfile.mkdtemp(prefix="openfhe_") @@ -547,7 +546,9 @@ def get_memory_usage(self): # --- Two-party key generation protocol --- # 1. Server (lead party) generates initial key pair print("Step 1: Server generates lead keys...") - server.openfhe_cc = OpenFHEThresholdCKKS(security_level=128, ring_dim=16384) + server.openfhe_cc = OpenFHEThresholdCKKS( + security_level=128, ring_dim=16384 + ) kp1 = server.openfhe_cc.generate_lead_keys() # Serialize context and lead public key to shared files @@ -562,9 +563,11 @@ def get_memory_usage(self): # Trainer reads serialized context + lead PK from files, generates its share pk2_path = os.path.join(he_dir, "pk_nonlead.bin") - ray.get(designated_trainer.setup_openfhe_nonlead.remote( - cc_path, pk1_path, pk2_path - )) + ray.get( + designated_trainer.setup_openfhe_nonlead.remote( + cc_path, pk1_path, pk2_path + ) + ) # 3. Server reads the non-lead public key (= joint PK) and uses it print("Step 3: Server reads joint public key...") @@ -624,6 +627,7 @@ def get_memory_usage(self): if use_lowrank: # Low-rank mode: decrypt each trainer separately, decompress, then add from fedgraph.low_rank.compression_utils import svd_decompress + u_shape = meta["U_shape"] s_len = meta["S_len"] v_shape = meta["V_shape"] @@ -631,7 +635,9 @@ def get_memory_usage(self): u_size = u_shape[0] * u_shape[1] v_size = v_shape[0] * v_shape[1] - print(f"Decrypting {len(ct_paths)} trainers x {num_chunks} chunks (low-rank, rank={meta['rank']})...") + print( + f"Decrypting {len(ct_paths)} trainers x {num_chunks} chunks (low-rank, rank={meta['rank']})..." + ) decrypted_tensor = torch.zeros(original_shape, dtype=torch.float32) for trainer_idx in range(len(ct_paths)): @@ -640,23 +646,31 @@ def get_memory_usage(self): partial_leads = [] for chunk_idx in range(num_chunks): chunk_path = f"{base}_chunk{chunk_idx}{ext}" - ct, ok = openfhe.DeserializeCiphertext(chunk_path, openfhe.BINARY) + ct, ok = openfhe.DeserializeCiphertext( + chunk_path, openfhe.BINARY + ) if not ok: - raise RuntimeError(f"Failed to deserialize {chunk_path}") + raise RuntimeError( + f"Failed to deserialize {chunk_path}" + ) partial_leads.append(server.openfhe_cc.partial_decrypt(ct)) # Write ct for designated trainer agg_path = os.path.join(he_dir, f"agg_ct_{chunk_idx}.bin") openfhe.SerializeToFile(agg_path, ct, openfhe.BINARY) # Batch partial decrypt on designated trainer - ray.get(designated_trainer.openfhe_partial_decrypt_main_batch.remote( - he_dir, num_chunks - )) + ray.get( + designated_trainer.openfhe_partial_decrypt_main_batch.remote( + he_dir, num_chunks + ) + ) # Fuse and collect values trainer_values = [] for chunk_idx in range(num_chunks): - partial_main_path = os.path.join(he_dir, f"partial_main_{chunk_idx}.bin") + partial_main_path = os.path.join( + he_dir, f"partial_main_{chunk_idx}.bin" + ) partial_main_ct, ok = openfhe.DeserializeCiphertext( partial_main_path, openfhe.BINARY ) @@ -669,43 +683,68 @@ def get_memory_usage(self): # Decompress SVD and accumulate vals = trainer_values[:total_elements] - U = torch.tensor(vals[:u_size], dtype=torch.float32).reshape(u_shape) - S = torch.tensor(vals[u_size:u_size + s_len], dtype=torch.float32) - V = torch.tensor(vals[u_size + s_len:u_size + s_len + v_size], dtype=torch.float32).reshape(v_shape) + U = torch.tensor(vals[:u_size], dtype=torch.float32).reshape( + u_shape + ) + S = torch.tensor( + vals[u_size : u_size + s_len], dtype=torch.float32 + ) + V = torch.tensor( + vals[u_size + s_len : u_size + s_len + v_size], + dtype=torch.float32, + ).reshape(v_shape) decrypted_tensor += svd_decompress(U, S, V) shape = original_shape - print(f"Low-rank decompressed and aggregated: rank={meta['rank']}, shape={shape}") + print( + f"Low-rank decompressed and aggregated: rank={meta['rank']}, shape={shape}" + ) else: # Standard mode: aggregate ciphertexts homomorphically, then decrypt - print(f"Aggregating {num_chunks} chunks across {len(ct_paths)} trainers...") + print( + f"Aggregating {num_chunks} chunks across {len(ct_paths)} trainers..." + ) partial_leads = [] for chunk_idx in range(num_chunks): agg_ct = None for i in range(len(ct_paths)): base, ext = os.path.splitext(ct_paths[i]) chunk_path = f"{base}_chunk{chunk_idx}{ext}" - ct, ok = openfhe.DeserializeCiphertext(chunk_path, openfhe.BINARY) + ct, ok = openfhe.DeserializeCiphertext( + chunk_path, openfhe.BINARY + ) if not ok: - raise RuntimeError(f"Failed to deserialize {chunk_path}") - agg_ct = ct if agg_ct is None else server.openfhe_cc.add_ciphertexts(agg_ct, ct) + raise RuntimeError( + f"Failed to deserialize {chunk_path}" + ) + agg_ct = ( + ct + if agg_ct is None + else server.openfhe_cc.add_ciphertexts(agg_ct, ct) + ) partial_leads.append(server.openfhe_cc.partial_decrypt(agg_ct)) agg_ct_path = os.path.join(he_dir, f"agg_ct_{chunk_idx}.bin") openfhe.SerializeToFile(agg_ct_path, agg_ct, openfhe.BINARY) print(f"Batch partial decryption on designated trainer...") - ray.get(designated_trainer.openfhe_partial_decrypt_main_batch.remote( - he_dir, num_chunks - )) + ray.get( + designated_trainer.openfhe_partial_decrypt_main_batch.remote( + he_dir, num_chunks + ) + ) all_fused_values = [] for chunk_idx in range(num_chunks): - partial_main_path = os.path.join(he_dir, f"partial_main_{chunk_idx}.bin") + partial_main_path = os.path.join( + he_dir, f"partial_main_{chunk_idx}.bin" + ) partial_main_ct, ok = openfhe.DeserializeCiphertext( partial_main_path, openfhe.BINARY ) if not ok: - raise RuntimeError(f"Failed to deserialize partial_main chunk {chunk_idx}") + raise RuntimeError( + f"Failed to deserialize partial_main chunk {chunk_idx}" + ) fused = server.openfhe_cc.cc.MultipartyDecryptFusion( [partial_leads[chunk_idx], partial_main_ct] ) @@ -729,11 +768,14 @@ def get_memory_usage(self): # Cleanup temp files import shutil + shutil.rmtree(he_dir, ignore_errors=True) pretrain_time = time.time() - pretrain_start pretrain_upload = sum(enc_sizes) / (1024 * 1024) # MB - pretrain_download = agg_size * len(server.trainers) / (1024 * 1024) # MB + pretrain_download = ( + agg_size * len(server.trainers) / (1024 * 1024) + ) # MB pretrain_comm_cost = pretrain_upload + pretrain_download # Print performance metrics @@ -742,9 +784,13 @@ def get_memory_usage(self): print(f"Aggregation Time: {aggregation_time:.2f} seconds") print(f"Pre-training Upload: {pretrain_upload:.2f} MB") print(f"Pre-training Download: {pretrain_download:.2f} MB") - print(f"Total Pre-training Communication Cost: {pretrain_comm_cost:.2f} MB") + print( + f"Total Pre-training Communication Cost: {pretrain_comm_cost:.2f} MB" + ) else: - raise ValueError(f"Unknown he_backend: {getattr(args, 'he_backend', None)}") + raise ValueError( + f"Unknown he_backend: {getattr(args, 'he_backend', None)}" + ) pretrain_time = time.time() - pretrain_start pretrain_upload = sum(enc_sizes) / (1024 * 1024) # MB pretrain_download = sum(download_sizes) / (1024 * 1024) # MB @@ -1212,7 +1258,7 @@ def __init__(self, *args: Any, **kwds: Any): # Create trainers (same as original) if args.use_huggingface: trainers = [ - Trainer.remote( + Trainer.remote( # type: ignore[attr-defined] rank=i, args_hidden=args_hidden, device=device, @@ -1222,7 +1268,7 @@ def __init__(self, *args: Any, **kwds: Any): ] else: trainers = [ - Trainer.remote( + Trainer.remote( # type: ignore[attr-defined] rank=i, args_hidden=args_hidden, device=device, @@ -1474,7 +1520,7 @@ def __init__(self, *args: Any, **kwds: Any): # Create trainers if args.use_huggingface: trainers = [ - Trainer.remote( + Trainer.remote( # type: ignore[attr-defined] rank=i, args_hidden=args_hidden, device=device, @@ -1484,7 +1530,7 @@ def __init__(self, *args: Any, **kwds: Any): ] else: trainers = [ - Trainer.remote( + Trainer.remote( # type: ignore[attr-defined] rank=i, args_hidden=args_hidden, device=device, diff --git a/fedgraph/gnn_models.py b/fedgraph/gnn_models.py index b07f0c5..fdb1bcc 100644 --- a/fedgraph/gnn_models.py +++ b/fedgraph/gnn_models.py @@ -519,7 +519,7 @@ def forward(self, data: torch_geometric.data.Data) -> torch.Tensor: The prediction of the model. """ x, edge_index, batch = data.x, data.edge_index, data.batch - x = self.pre(x) + x = self.pre(x) # type: ignore[misc] for i in range(len(self.graph_convs)): x = self.graph_convs[i](x, edge_index) x = F.relu(x) @@ -527,7 +527,7 @@ def forward(self, data: torch_geometric.data.Data) -> torch.Tensor: x = global_add_pool(x, batch) x = self.post(x) x = F.dropout(x, self.dropout, training=self.training) - x = self.readout(x) + x = self.readout(x) # type: ignore[misc] x = F.log_softmax(x, dim=1) return x diff --git a/fedgraph/low_rank/server_lowrank.py b/fedgraph/low_rank/server_lowrank.py index eefadc6..435cd79 100644 --- a/fedgraph/low_rank/server_lowrank.py +++ b/fedgraph/low_rank/server_lowrank.py @@ -33,7 +33,7 @@ def __init__( self.energy_threshold = getattr(args, "energy_threshold", 0.95) self.fixed_rank = getattr(args, "fixed_rank", 10) - self.compression_stats = [] + self.compression_stats: list[Dict[str, Any]] = [] print(f"Server initialized with low-rank compression: {self.use_lowrank}") if self.use_lowrank: diff --git a/fedgraph/low_rank/trainer_lowrank.py b/fedgraph/low_rank/trainer_lowrank.py index c32bc64..e9bce6b 100644 --- a/fedgraph/low_rank/trainer_lowrank.py +++ b/fedgraph/low_rank/trainer_lowrank.py @@ -20,11 +20,11 @@ def get_compressed_params(self) -> Dict[str, Any]: Get model parameters with optional compression. """ if not self.use_lowrank: - return {"params": dict(self.model.named_parameters()), "compressed": False} + return {"params": dict(self.model.named_parameters()), "compressed": False} # type: ignore[attr-defined] params = { name: param.data.cpu().detach() - for name, param in self.model.named_parameters() + for name, param in self.model.named_parameters() # type: ignore[attr-defined] } compressed_params = {} @@ -55,18 +55,18 @@ def update_compressed_params( if not compressed_data.get("compressed", False): # Standard parameter update params = compressed_data["params"] - self.model.to("cpu") - for name, param in self.model.named_parameters(): + self.model.to("cpu") # type: ignore[attr-defined] + for name, param in self.model.named_parameters(): # type: ignore[attr-defined] if name in params: param.data.copy_(params[name]) - self.model.to(self.device) + self.model.to(self.device) # type: ignore[attr-defined] return # Decompress and update - self.model.to("cpu") + self.model.to("cpu") # type: ignore[attr-defined] compressed_params = compressed_data["params"] - for name, param in self.model.named_parameters(): + for name, param in self.model.named_parameters(): # type: ignore[attr-defined] if name in compressed_params: param_data = compressed_params[name] if isinstance(param_data, dict) and "U" in param_data: @@ -79,4 +79,4 @@ def update_compressed_params( # Direct copy param.data.copy_(param_data) - self.model.to(self.device) + self.model.to(self.device) # type: ignore[attr-defined] diff --git a/fedgraph/openfhe_threshold.py b/fedgraph/openfhe_threshold.py index d9a44e9..3ecbea5 100644 --- a/fedgraph/openfhe_threshold.py +++ b/fedgraph/openfhe_threshold.py @@ -13,12 +13,13 @@ - Fusion combines both partial decryptions """ -import openfhe -import numpy as np -from typing import List, Tuple, Optional, Union import logging -import tempfile import os +import tempfile +from typing import List, Optional, Tuple, Union + +import numpy as np +import openfhe logger = logging.getLogger(__name__) @@ -63,7 +64,9 @@ def _setup_context(self): self.cc.Enable(openfhe.ADVANCEDSHE) self.cc.Enable(openfhe.MULTIPARTY) - logger.info(f"OpenFHE context initialized (ring_dim={self.cc.GetRingDimension()})") + logger.info( + f"OpenFHE context initialized (ring_dim={self.cc.GetRingDimension()})" + ) def generate_lead_keys(self): """Lead party (server): generate initial key pair.""" @@ -166,7 +169,9 @@ def serialize_public_key(self) -> bytes: tmpdir = tempfile.mkdtemp() pk_path = os.path.join(tmpdir, "pk.bin") try: - self.cc.SerializeToFile(pk_path, openfhe.BINARY) # context must be serialized first + self.cc.SerializeToFile( + pk_path, openfhe.BINARY + ) # context must be serialized first if not openfhe.SerializeToFile(pk_path, self.public_key, openfhe.BINARY): raise RuntimeError("Failed to serialize public key") with open(pk_path, "rb") as f: @@ -217,7 +222,9 @@ def get_context_info(self) -> dict: } -def create_threshold_context(security_level: int = 128, ring_dim: int = 16384) -> OpenFHEThresholdCKKS: +def create_threshold_context( + security_level: int = 128, ring_dim: int = 16384 +) -> OpenFHEThresholdCKKS: """Create a new threshold HE context.""" return OpenFHEThresholdCKKS(security_level, ring_dim) @@ -251,7 +258,11 @@ def timeout_handler(signum, frame): # 3) Server also sets the joint public key (kp2.publicKey) server.set_public_key(kp2.publicKey) - print("Joint PK set on both?", server.public_key is not None, trainer.public_key is not None) + print( + "Joint PK set on both?", + server.public_key is not None, + trainer.public_key is not None, + ) print("Lead/Main flags:", server.is_lead_party, trainer.is_lead_party) # Encrypt test vectors diff --git a/fedgraph/server_class.py b/fedgraph/server_class.py index 3b2b635..d17b583 100644 --- a/fedgraph/server_class.py +++ b/fedgraph/server_class.py @@ -15,7 +15,7 @@ try: from fedgraph.openfhe_threshold import OpenFHEThresholdCKKS # noqa: F401 except ImportError: # pragma: no cover - OpenFHEThresholdCKKS = None # type: ignore[assignment] + OpenFHEThresholdCKKS = None # type: ignore[misc,assignment] from dtaidistance import dtw from fedgraph.gnn_models import ( @@ -126,7 +126,7 @@ def __init__( self.he_backend = getattr(args, "he_backend", "tenseal") if self.he_backend == "openfhe": self.openfhe_cc = OpenFHEThresholdCKKS() - self.aggregation_stats = [] + self.aggregation_stats: list[Any] = [] print("Initialized OpenFHE threshold context") else: self.he_backend = "tenseal" @@ -183,7 +183,9 @@ def mask_encrypted_feature_sum(self, encrypted_feature_sum, node_indexes): selected_nodes.min().item() < 0 or selected_nodes.max().item() >= num_nodes ): - raise ValueError("node_indexes contains values outside feature sum shape") + raise ValueError( + "node_indexes contains values outside feature sum shape" + ) mask = torch.zeros((num_nodes, feature_dim), dtype=torch.float64) mask[selected_nodes] = 1.0 @@ -201,7 +203,9 @@ def aggregate_encrypted_params(self, encrypted_params_list): first_params, _ = encrypted_params_list[0] n_layers = len(first_params) if n_layers == 0: - raise ValueError("encrypted parameter payload must contain at least one layer") + raise ValueError( + "encrypted parameter payload must contain at least one layer" + ) layer_shapes = [] layer_scale_values = [[] for _ in range(n_layers)] @@ -213,7 +217,9 @@ def aggregate_encrypted_params(self, encrypted_params_list): if len(trainer_params) != n_layers: raise ValueError("all trainers must provide the same number of layers") if len(trainer_metadata) != n_layers: - raise ValueError("metadata length must match encrypted parameter layers") + raise ValueError( + "metadata length must match encrypted parameter layers" + ) validated_metadata = [] for layer_idx, metadata in enumerate(trainer_metadata): @@ -244,7 +250,10 @@ def aggregate_encrypted_params(self, encrypted_params_list): validated_params.append((trainer_params, validated_metadata)) aggregation_metadata = [ - {"shape": layer_shapes[layer_idx], "scale": max(layer_scale_values[layer_idx])} + { + "shape": layer_shapes[layer_idx], + "scale": max(layer_scale_values[layer_idx]), + } for layer_idx in range(n_layers) ] diff --git a/fedgraph/trainer_class.py b/fedgraph/trainer_class.py index b3c3721..9ff9564 100644 --- a/fedgraph/trainer_class.py +++ b/fedgraph/trainer_class.py @@ -4,7 +4,7 @@ import time import warnings from io import BytesIO -from typing import Any, Dict, List, Union +from typing import Any, Dict, List, Optional, Union logging.basicConfig(level=logging.INFO) # logger = logging.getLogger(__name__) @@ -30,13 +30,15 @@ GCN_arxiv, SAGE_products, ) + # Threshold-HE backend is optional. We delay-import OpenFHE bindings here so the # rest of fedgraph stays importable on systems without the OpenFHE wheel. try: from fedgraph.openfhe_threshold import OpenFHEThresholdCKKS # noqa: F401 + _OPENFHE_AVAILABLE = True except ImportError: # pragma: no cover - exercised only when openfhe is missing - OpenFHEThresholdCKKS = None # type: ignore[assignment] + OpenFHEThresholdCKKS = None # type: ignore[misc,assignment] _OPENFHE_AVAILABLE = False from fedgraph.train_func import test, train @@ -162,8 +164,8 @@ def __init__( features: torch.Tensor = None, idx_train: torch.Tensor = None, idx_test: torch.Tensor = None, - global_node_num: int = None, - class_num: int = None, + global_node_num: Optional[int] = None, + class_num: Optional[int] = None, ): # from gnn_models import GCN_Graph_Classification # Per-trainer seed = global_seed * 1000 + rank (lets us vary across runs @@ -296,7 +298,9 @@ def init_model(self, global_node_num, class_num): dropout=0.5, NumLayers=self.args.num_layers, ).to(self.device) - elif "ogbn" in self.args.dataset: # ogbn large datasets default to GCN_arxiv + elif ( + "ogbn" in self.args.dataset + ): # ogbn large datasets default to GCN_arxiv print("Running GCN_arxiv") self.model = GCN_arxiv( nfeat=self.features.shape[1], @@ -332,7 +336,7 @@ def update_params(self, params: tuple, current_global_epoch: int) -> None: # load global parameter from global server if self.model is None: return - + self.model.to("cpu") for ( p, @@ -384,10 +388,12 @@ def get_local_feature_sum(self) -> torch.Tensor: # Sum of features of all 1-hop nodes for each node one_hop_neighbor_feature_sum = get_1hop_feature_sum( - new_feature_for_trainer, self.adj, self.device, - norm_type=getattr(self.args, "norm_type", "none") + new_feature_for_trainer, + self.adj, + self.device, + norm_type=getattr(self.args, "norm_type", "none"), ) - if hasattr(self.args, 'use_encryption') and self.args.use_encryption: + if hasattr(self.args, "use_encryption") and self.args.use_encryption: print( f"Trainer {self.rank} - Original feature sum (first 10 and last 10 elements): " f"{one_hop_neighbor_feature_sum.flatten()[:10].tolist()} ... {one_hop_neighbor_feature_sum.flatten()[-10:].tolist()}" @@ -411,8 +417,10 @@ def get_local_feature_sum_og(self) -> torch.Tensor: ).to(self.device) new_feature_for_trainer[self.local_node_index] = self.features one_hop_neighbor_feature_sum = get_1hop_feature_sum( - new_feature_for_trainer, self.adj, self.device, - norm_type=getattr(self.args, "norm_type", "none") + new_feature_for_trainer, + self.adj, + self.device, + norm_type=getattr(self.args, "norm_type", "none"), ) computation_time = time.time() - computation_start @@ -473,8 +481,8 @@ def decrypt_feature_sum(self, encrypted_sum, shape): def get_encrypted_local_feature_sum(self, ct_output_path=None): # Check HE backend and route accordingly - if hasattr(self, 'he_backend') and self.he_backend == "openfhe": - if getattr(self, 'use_lowrank', False): + if hasattr(self, "he_backend") and self.he_backend == "openfhe": + if getattr(self, "use_lowrank", False): return self._get_openfhe_lowrank_encrypted_feature_sum(ct_output_path) return self._get_openfhe_encrypted_local_feature_sum(ct_output_path) else: @@ -482,19 +490,22 @@ def get_encrypted_local_feature_sum(self, ct_output_path=None): def _get_openfhe_encrypted_local_feature_sum(self, ct_output_path=None): """OpenFHE encryption of local feature sum, chunked and serialized to files.""" - import openfhe import json + import openfhe + new_feature_for_trainer = torch.zeros( self.global_node_num, self.features.shape[1] ).to(self.device) new_feature_for_trainer[self.local_node_index] = self.features feature_sum = get_1hop_feature_sum( - new_feature_for_trainer, self.adj, self.device, - norm_type=getattr(self.args, "norm_type", "none") + new_feature_for_trainer, + self.adj, + self.device, + norm_type=getattr(self.args, "norm_type", "none"), ) - if not hasattr(self, 'openfhe_cc'): + if not hasattr(self, "openfhe_cc"): raise RuntimeError("OpenFHE context not available on trainer") encryption_start = time.time() @@ -516,16 +527,24 @@ def _get_openfhe_encrypted_local_feature_sum(self, ct_output_path=None): # Write metadata meta_path = f"{base}_meta.json" with open(meta_path, "w") as f: - json.dump({"num_chunks": num_chunks, "slot_count": slot_count, - "total_elements": len(feature_list)}, f) + json.dump( + { + "num_chunks": num_chunks, + "slot_count": slot_count, + "total_elements": len(feature_list), + }, + f, + ) encryption_time = time.time() - encryption_start return feature_sum.shape, encryption_time def _get_openfhe_lowrank_encrypted_feature_sum(self, ct_output_path=None): """Low-rank compress feature sum, then encrypt with OpenFHE threshold HE.""" - import openfhe import json + + import openfhe + from fedgraph.low_rank.compression_utils import svd_compress new_feature_for_trainer = torch.zeros( @@ -533,11 +552,13 @@ def _get_openfhe_lowrank_encrypted_feature_sum(self, ct_output_path=None): ).to(self.device) new_feature_for_trainer[self.local_node_index] = self.features feature_sum = get_1hop_feature_sum( - new_feature_for_trainer, self.adj, self.device, - norm_type=getattr(self.args, "norm_type", "none") + new_feature_for_trainer, + self.adj, + self.device, + norm_type=getattr(self.args, "norm_type", "none"), ) - if not hasattr(self, 'openfhe_cc'): + if not hasattr(self, "openfhe_cc"): raise RuntimeError("OpenFHE context not available on trainer") encryption_start = time.time() @@ -568,17 +589,20 @@ def _get_openfhe_lowrank_encrypted_feature_sum(self, ct_output_path=None): # Write metadata including SVD shape info for reconstruction meta_path = f"{base}_meta.json" with open(meta_path, "w") as f: - json.dump({ - "num_chunks": num_chunks, - "slot_count": slot_count, - "total_elements": len(all_values), - "lowrank": True, - "rank": rank, - "U_shape": list(U.shape), - "S_len": len(s_flat), - "V_shape": list(V.shape), - "original_shape": list(feature_sum.shape), - }, f) + json.dump( + { + "num_chunks": num_chunks, + "slot_count": slot_count, + "total_elements": len(all_values), + "lowrank": True, + "rank": rank, + "U_shape": list(U.shape), + "S_len": len(s_flat), + "V_shape": list(V.shape), + "original_shape": list(feature_sum.shape), + }, + f, + ) encryption_time = time.time() - encryption_start return feature_sum.shape, encryption_time @@ -591,8 +615,10 @@ def _get_tenseal_encrypted_local_feature_sum(self): ).to(self.device) new_feature_for_trainer[self.local_node_index] = self.features feature_sum = get_1hop_feature_sum( - new_feature_for_trainer, self.adj, self.device, - norm_type=getattr(self.args, "norm_type", "none") + new_feature_for_trainer, + self.adj, + self.device, + norm_type=getattr(self.args, "norm_type", "none"), ) # Encrypt the feature sum @@ -606,6 +632,7 @@ def _get_tenseal_encrypted_local_feature_sum(self): def setup_openfhe_nonlead(self, cc_path, lead_pk_path, output_pk_path): """Setup OpenFHE as non-lead party using file-based serialization.""" import openfhe + from fedgraph.openfhe_threshold import OpenFHEThresholdCKKS # Deserialize context from file @@ -642,6 +669,7 @@ def setup_openfhe_nonlead(self, cc_path, lead_pk_path, output_pk_path): def set_openfhe_public_key(self, cc_path, joint_pk_path): """Set the joint public key for encryption-only trainers using file-based serialization.""" import openfhe + from fedgraph.openfhe_threshold import OpenFHEThresholdCKKS # Deserialize context @@ -671,7 +699,7 @@ def openfhe_partial_decrypt_main_batch(self, he_dir, num_chunks): """Batch partial decryption of all chunks at once.""" import openfhe - if not hasattr(self, 'openfhe_cc'): + if not hasattr(self, "openfhe_cc"): raise RuntimeError("OpenFHE context not initialized on trainer") for chunk_idx in range(num_chunks): @@ -687,7 +715,9 @@ def openfhe_partial_decrypt_main_batch(self, he_dir, num_chunks): ) openfhe.SerializeToFile(partial_path, partial_list[0], openfhe.BINARY) - print(f"Trainer {self.rank}: Batch partial decryption done ({num_chunks} chunks)") + print( + f"Trainer {self.rank}: Batch partial decryption done ({num_chunks} chunks)" + ) return True def load_encrypted_feature_aggregation(self, encrypted_data): @@ -737,18 +767,18 @@ def load_encrypted_params(self, encrypted_data: tuple, current_global_epoch: int """Load encrypted parameters with rescaling""" params_list, metadata = encrypted_data - self.model.to("cpu") + self.model.to("cpu") # type: ignore[attr-defined] # load each layer's parameters for param, enc_param, meta in zip( - self.model.parameters(), params_list, metadata + self.model.parameters(), params_list, metadata # type: ignore[attr-defined] ): - decrypted = ts.ckks_vector_from(self.he_context, enc_param).decrypt() + decrypted = ts.ckks_vector_from(self.he_context, enc_param).decrypt() # type: ignore[attr-defined] param_data = torch.tensor(decrypted).reshape(meta["shape"]) param_data = param_data / meta["scale"] # Reverse scaling param.data.copy_(param_data) - self.model.to(self.device) + self.model.to(self.device) # type: ignore[attr-defined] return True def use_fedavg_feature(self) -> None: @@ -908,7 +938,11 @@ def local_test(self) -> list: test_labels = self.test_labels.to(self.device) idx_test = self.idx_test.to(self.device) local_test_loss, local_test_acc = test( - self.model, feats, adj, test_labels, idx_test, + self.model, + feats, + adj, + test_labels, + idx_test, ) self.test_losses.append(local_test_loss) self.test_accs.append(local_test_acc) @@ -1649,8 +1683,10 @@ def get_train_test_data_at_current_time_step( buffer_size : int, optional The size of the buffer. The default is 10. """ - print("loading buffer_train_data_list") if use_buffer else print( - "loading train_data and test_data" + ( + print("loading buffer_train_data_list") + if use_buffer + else print("loading train_data and test_data") ) load_res = get_data_loaders_per_time_step( diff --git a/fedgraph/utils_lp.py b/fedgraph/utils_lp.py index 1bc22d0..02c898d 100644 --- a/fedgraph/utils_lp.py +++ b/fedgraph/utils_lp.py @@ -240,14 +240,14 @@ def get_data_by_time_step( ] # only keeps the edges within the time range if constant_edges != -1: # only keep the last `constant_edges` edges - data_current_time_step[ - "user", "select", "item" - ].edge_index = data_current_time_step["user", "select", "item"].edge_index[ - :, -constant_edges: - ] - data_current_time_step[ - "user", "select", "item" - ].edge_attr = None # no need to store the edge attributes + data_current_time_step["user", "select", "item"].edge_index = ( + data_current_time_step["user", "select", "item"].edge_index[ + :, -constant_edges: + ] + ) + data_current_time_step["user", "select", "item"].edge_attr = ( + None # no need to store the edge attributes + ) data_current_time_step = torch_geometric.transforms.ToUndirected()( data_current_time_step diff --git a/fedgraph/utils_nc.py b/fedgraph/utils_nc.py index de4d24d..eba0adc 100644 --- a/fedgraph/utils_nc.py +++ b/fedgraph/utils_nc.py @@ -123,7 +123,7 @@ def label_dirichlet_partition( max_attempts = 1000 # increased // can revoke later while attempts < max_attempts: attempts += 1 - idx_batch = [[] for _ in range(n_parties)] + idx_batch: list[list[int]] = [[] for _ in range(n_parties)] for k in range(K): idx_k = label_indices[k] @@ -443,15 +443,17 @@ def get_1hop_feature_sum( if norm_type == "sym": # Symmetric normalization: D^{-1/2} A D^{-1/2} deg_inv_sqrt = deg.pow(-0.5) - deg_inv_sqrt[deg_inv_sqrt == float('inf')] = 0 + deg_inv_sqrt[deg_inv_sqrt == float("inf")] = 0 edge_weight = deg_inv_sqrt[src] * deg_inv_sqrt[dst] elif norm_type == "row": # Row normalization (mean aggregation): D^{-1} A deg_inv = deg.pow(-1) - deg_inv[deg_inv == float('inf')] = 0 + deg_inv[deg_inv == float("inf")] = 0 edge_weight = deg_inv[src] else: - raise ValueError(f"Unknown norm_type: {norm_type}. Use 'sym', 'row', or 'none'.") + raise ValueError( + f"Unknown norm_type: {norm_type}. Use 'sym', 'row', or 'none'." + ) adjacency_matrix = torch.sparse_coo_tensor( edge_with_self, diff --git a/pyproject.toml b/pyproject.toml index 8bce10d..20367c3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,3 +5,7 @@ build-backend = "setuptools.build_meta" # third party tools [tool.isort] profile = "black" + +[tool.mypy] +python_version = "3.9" +ignore_missing_imports = true diff --git a/run_docker_openfhe.sh b/run_docker_openfhe.sh index 0e9a804..cb43735 100755 --- a/run_docker_openfhe.sh +++ b/run_docker_openfhe.sh @@ -24,4 +24,4 @@ docker run -it --rm \ fedgraph-openfhe \ /bin/bash -echo "👋 Container stopped." \ No newline at end of file +echo "👋 Container stopped." diff --git a/setup.py b/setup.py index 23022d2..a4b8121 100644 --- a/setup.py +++ b/setup.py @@ -54,7 +54,7 @@ "ogb", ], extras_require={ - "dev": ["build", "mypy", "pre-commit", "pytest"], + "dev": ["build", "mypy>=1.20,<2", "pre-commit>=4.6,<5", "pytest"], # Optional threshold-HE backend used by FedGCN-v2. Install with # pip install "fedgraph[openfhe]" # OpenFHE Python wheels are only published for Linux/manylinux at the diff --git a/test_openfhe_smoke.py b/test_openfhe_smoke.py index 45575e6..ec4c5a1 100644 --- a/test_openfhe_smoke.py +++ b/test_openfhe_smoke.py @@ -2,9 +2,12 @@ """ Minimal OpenFHE CKKS smoke test to verify setup before trying multiparty. """ -import openfhe + import time +import openfhe + + def test_basic_ckks(): """Test basic CKKS functionality (no multiparty).""" print("🔍 Testing basic OpenFHE CKKS...") @@ -51,7 +54,7 @@ def test_basic_ckks(): print(f"Result: {result[:len(x)]}") # Verify accuracy - errors = [abs(e - r) for e, r in zip(x, result[:len(x)])] + errors = [abs(e - r) for e, r in zip(x, result[: len(x)])] max_error = max(errors) print(f"Max error: {max_error:.2e}") @@ -62,11 +65,13 @@ def test_basic_ckks(): print("❌ Basic CKKS test FAILED!") return False + def test_import_speed(): """Test OpenFHE import speed.""" print("🔍 Testing OpenFHE import speed...") start = time.time() import openfhe + import_time = time.time() - start print(f"✅ Import took {import_time:.2f} seconds") @@ -77,6 +82,7 @@ def test_import_speed(): print("⚠️ Import is slow (possible emulation)") return False + if __name__ == "__main__": print("🚀 OpenFHE Smoke Test") print("=" * 50) diff --git a/tests/conftest.py b/tests/conftest.py index 4c4191b..e916697 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -7,6 +7,7 @@ * Skip markers for tests that require optional encryption backends (``needs_openfhe``, ``needs_tenseal``). """ + from __future__ import annotations import os @@ -19,14 +20,15 @@ import torch from torch_geometric.data import Data - # --------------------------------------------------------------------------- # Optional encryption backends -- skip markers # --------------------------------------------------------------------------- + def _openfhe_available() -> bool: try: import openfhe # noqa: F401 + return True except Exception: # pragma: no cover - exercised only on systems without openfhe return False @@ -35,6 +37,7 @@ def _openfhe_available() -> bool: def _tenseal_available() -> bool: try: import tenseal # noqa: F401 + return True except Exception: # pragma: no cover return False @@ -55,6 +58,7 @@ def _tenseal_available() -> bool: # Common fixtures # --------------------------------------------------------------------------- + @pytest.fixture def temp_dir(): """Create a temporary directory for test files.""" @@ -74,7 +78,9 @@ def sample_graph_data(): x = torch.randn(num_nodes, num_features) # Create edges (simple chain graph) - edge_index = torch.tensor([[i, i + 1] for i in range(num_nodes - 1)]).t().contiguous() + edge_index = ( + torch.tensor([[i, i + 1] for i in range(num_nodes - 1)]).t().contiguous() + ) # Create labels y = torch.randint(0, num_classes, (num_nodes,)) @@ -112,11 +118,9 @@ def mock_args(): @pytest.fixture def mock_ray_cluster(): """Mock Ray cluster for testing.""" - with patch("ray.init") as mock_init, \ - patch("ray.get") as mock_get, \ - patch("ray.put") as mock_put, \ - patch("ray.remote") as mock_remote: - + with patch("ray.init") as mock_init, patch("ray.get") as mock_get, patch( + "ray.put" + ) as mock_put, patch("ray.remote") as mock_remote: mock_init.return_value = None mock_get.side_effect = lambda x: x mock_put.side_effect = lambda x: x diff --git a/tests/integration/test_fedgraph_integration.py b/tests/integration/test_fedgraph_integration.py index 300422c..75dbbf6 100644 --- a/tests/integration/test_fedgraph_integration.py +++ b/tests/integration/test_fedgraph_integration.py @@ -1,9 +1,10 @@ +import os +import tempfile +from unittest.mock import MagicMock, Mock, patch + +import numpy as np import pytest import torch -import numpy as np -from unittest.mock import Mock, patch, MagicMock -import tempfile -import os """ Integration tests for FedGraph library. @@ -17,53 +18,63 @@ @pytest.mark.integration class TestFedGraphNCWorkflow: """Integration tests for Node Classification workflow.""" - + def setup_method(self): """Setup test data for NC integration tests.""" self.num_nodes = 100 self.num_features = 50 self.num_classes = 7 self.num_trainers = 3 - + # Create mock graph data self.features = torch.randn(self.num_nodes, self.num_features) self.labels = torch.randint(0, self.num_classes, (self.num_nodes,)) self.edge_index = torch.randint(0, self.num_nodes, (2, 200)) self.idx_train = torch.arange(0, 70) self.idx_test = torch.arange(85, 100) - + # Create node splits for trainers split_size = self.num_nodes // self.num_trainers self.split_node_indexes = [ torch.arange(i * split_size, min((i + 1) * split_size, self.num_nodes)) for i in range(self.num_trainers) ] - + # Handle remainder nodes if sum(len(split) for split in self.split_node_indexes) < self.num_nodes: remainder = torch.arange(split_size * self.num_trainers, self.num_nodes) - self.split_node_indexes[-1] = torch.cat([self.split_node_indexes[-1], remainder]) - - @patch('fedgraph.data_process.NC_load_data') - @patch('fedgraph.data_process.label_dirichlet_partition') - @patch('fedgraph.data_process.get_in_comm_indexes') - def test_nc_data_processing_pipeline(self, mock_get_indexes, mock_partition, mock_load_data): + self.split_node_indexes[-1] = torch.cat( + [self.split_node_indexes[-1], remainder] + ) + + @patch("fedgraph.data_process.NC_load_data") + @patch("fedgraph.data_process.label_dirichlet_partition") + @patch("fedgraph.data_process.get_in_comm_indexes") + def test_nc_data_processing_pipeline( + self, mock_get_indexes, mock_partition, mock_load_data + ): """Test the complete NC data processing pipeline.""" from fedgraph.data_process import data_loader_NC - + # Setup mocks mock_load_data.return_value = ( - self.features, Mock(), self.labels, - self.idx_train, Mock(), self.idx_test + self.features, + Mock(), + self.labels, + self.idx_train, + Mock(), + self.idx_test, ) - mock_partition.return_value = [split.tolist() for split in self.split_node_indexes] + mock_partition.return_value = [ + split.tolist() for split in self.split_node_indexes + ] mock_get_indexes.return_value = ( - {i: torch.arange(i*20, (i+1)*20) for i in range(self.num_trainers)}, + {i: torch.arange(i * 20, (i + 1) * 20) for i in range(self.num_trainers)}, {i: torch.arange(0, 10) for i in range(self.num_trainers)}, {i: torch.arange(10, 15) for i in range(self.num_trainers)}, - {i: self.edge_index for i in range(self.num_trainers)} + {i: self.edge_index for i in range(self.num_trainers)}, ) - + # Create mock args args = Mock() args.use_huggingface = False @@ -72,30 +83,30 @@ def test_nc_data_processing_pipeline(self, mock_get_indexes, mock_partition, moc args.iid_beta = 0.5 args.distribution_type = "dirichlet" args.num_hops = 2 - + # Test the pipeline result = data_loader_NC(args) - + assert len(result) == 11 # Expected return tuple length assert mock_load_data.called assert mock_partition.called assert mock_get_indexes.called - - @patch('fedgraph.trainer_class.GCN') + + @patch("fedgraph.trainer_class.GCN") def test_nc_trainer_initialization(self, mock_gcn_class): """Test NC trainer initialization and basic operations.""" from fedgraph.trainer_class import Trainer_General - + # Mock model mock_model = Mock() mock_gcn_class.return_value = mock_model mock_model.to.return_value = mock_model - + # Create trainer trainer = Trainer_General( rank=0, args_hidden=64, - device=torch.device('cpu'), + device=torch.device("cpu"), args=Mock(local_step=5, method="FedAvg"), local_node_index=self.split_node_indexes[0], communicate_node_index=torch.arange(0, 50), @@ -104,31 +115,31 @@ def test_nc_trainer_initialization(self, mock_gcn_class): test_labels=self.labels[30:40], features=self.features, idx_train=torch.arange(0, 20), - idx_test=torch.arange(20, 30) + idx_test=torch.arange(20, 30), ) - + # Test basic operations info = trainer.get_info() assert "features_num" in info assert "label_num" in info - + # Test parameter operations params = trainer.get_params() assert isinstance(params, tuple) - + rank = trainer.get_rank() assert rank == 0 - - @patch('fedgraph.server_class.AggreGCN') + + @patch("fedgraph.server_class.AggreGCN") def test_nc_server_aggregation(self, mock_aggre_gcn): """Test NC server aggregation workflow.""" from fedgraph.server_class import Server - + # Mock model mock_model = Mock() mock_aggre_gcn.return_value = mock_model mock_model.to.return_value = mock_model - + # Mock trainers trainers = [] for i in range(self.num_trainers): @@ -137,117 +148,111 @@ def test_nc_server_aggregation(self, mock_aggre_gcn): trainer.update_params = Mock() trainer.get_params.return_value = (torch.randn(64, 50), torch.randn(64)) trainers.append(trainer) - + # Create server args = Mock() args.num_hops = 2 args.dataset = "cora" args.num_layers = 2 - + server = Server( feature_dim=self.num_features, args_hidden=64, class_num=self.num_classes, - device=torch.device('cpu'), + device=torch.device("cpu"), trainers=trainers, - args=args + args=args, ) - + # Test server operations assert server.num_of_trainers == self.num_trainers - + # Test parameter broadcast - mock_model.state_dict.return_value = {'weight': torch.randn(64, 50)} + mock_model.state_dict.return_value = {"weight": torch.randn(64, 50)} mock_model.parameters.return_value = [torch.randn(64, 50)] with patch("fedgraph.server_class.ray.get"): server.broadcast_params(current_global_epoch=1) - + # Verify all trainers received updates for trainer in trainers: trainer.update_params.remote.assert_called() -@pytest.mark.integration +@pytest.mark.integration class TestFedGraphGCWorkflow: """Integration tests for Graph Classification workflow.""" - + def setup_method(self): """Setup test data for GC integration tests.""" self.num_trainers = 3 self.num_node_features = 10 self.num_graph_labels = 2 self.train_size = 100 - + # Mock graph classification data self.mock_data = {} for i in range(self.num_trainers): trainer_name = f"{i}-PROTEINS" self.mock_data[trainer_name] = ( - { - 'train': Mock(), - 'val': Mock(), - 'test': Mock() - }, + {"train": Mock(), "val": Mock(), "test": Mock()}, self.num_node_features, self.num_graph_labels, - self.train_size + self.train_size, ) - - @patch('fedgraph.utils_gc.setup_server') - @patch('fedgraph.utils_gc.setup_trainers') - @patch('fedgraph.federated_methods.run_GC_Fed_algorithm') - def test_gc_complete_workflow(self, mock_run_fed, mock_setup_trainers, mock_setup_server): + + @patch("fedgraph.utils_gc.setup_server") + @patch("fedgraph.utils_gc.setup_trainers") + @patch("fedgraph.federated_methods.run_GC_Fed_algorithm") + def test_gc_complete_workflow( + self, mock_run_fed, mock_setup_trainers, mock_setup_server + ): """Test complete GC federated learning workflow.""" from fedgraph.federated_methods import run_GC - + # Mock server and trainers mock_server = Mock() mock_setup_server.return_value = mock_server - + mock_trainers = [Mock() for _ in range(self.num_trainers)] mock_setup_trainers.return_value = mock_trainers - + # Create args args = Mock() args.dataset = "PROTEINS" args.method = "FedAvg" args.federated_method = "FedAvg" args.monitor = False - + # Run the workflow run_GC(args, self.mock_data) - + # Verify setup was called mock_setup_server.assert_called_once() mock_setup_trainers.assert_called_once() mock_run_fed.assert_called_once() - + def test_gc_trainer_creation(self): """Test GC trainer creation and initialization.""" from fedgraph.trainer_class import Trainer_GC - + # Mock model mock_model = Mock() mock_model.named_parameters.return_value = [ - ('layer1.weight', torch.randn(10, 5, requires_grad=True)), - ('layer1.bias', torch.randn(10, requires_grad=True)) + ("layer1.weight", torch.randn(10, 5, requires_grad=True)), + ("layer1.bias", torch.randn(10, requires_grad=True)), ] mock_model.to.return_value = mock_model - + # Mock dataloader - mock_dataloader = { - 'train': Mock(), - 'val': Mock(), - 'test': Mock() - } - + mock_dataloader = {"train": Mock(), "val": Mock(), "test": Mock()} + # Mock optimizer mock_optimizer = Mock() - + # Mock args args = Mock() - args.device = torch.device('cpu') - + args.device = torch.device("cpu") + # Create trainer trainer = Trainer_GC( model=mock_model, @@ -256,38 +261,39 @@ def test_gc_trainer_creation(self): train_size=100, dataloader=mock_dataloader, optimizer=mock_optimizer, - args=args + args=args, ) - + # Test trainer properties assert trainer.id == 0 assert trainer.name == "trainer_0" assert trainer.train_size == 100 - assert hasattr(trainer, 'W') - assert hasattr(trainer, 'dW') + assert hasattr(trainer, "W") + assert hasattr(trainer, "dW") @pytest.mark.integration class TestFedGraphLPWorkflow: """Integration tests for Link Prediction workflow.""" - - @patch('fedgraph.utils_lp.check_data_files_existance') - @patch('fedgraph.utils_lp.get_global_user_item_mapping') - @patch('fedgraph.utils_lp.get_start_end_time') - @patch('fedgraph.federated_methods.Server_LP') - def test_lp_basic_setup(self, mock_server_lp, mock_get_time, - mock_get_mapping, mock_check_files): + + @patch("fedgraph.utils_lp.check_data_files_existance") + @patch("fedgraph.utils_lp.get_global_user_item_mapping") + @patch("fedgraph.utils_lp.get_start_end_time") + @patch("fedgraph.federated_methods.Server_LP") + def test_lp_basic_setup( + self, mock_server_lp, mock_get_time, mock_get_mapping, mock_check_files + ): """Test basic LP workflow setup.""" from fedgraph.federated_methods import run_LP - + # Setup mocks mock_check_files.return_value = True mock_get_mapping.return_value = ({1: 0, 2: 1}, {10: 0, 20: 1}) mock_get_time.return_value = (1000, 5000) - + mock_server = Mock() mock_server_lp.return_value = mock_server - + # Create args args = Mock() args.dataset = "ml-1m" @@ -295,14 +301,14 @@ def test_lp_basic_setup(self, mock_server_lp, mock_get_time, args.num_trainer = 3 args.device = "cpu" args.monitor = False - - with patch('fedgraph.federated_methods.ray'): + + with patch("fedgraph.federated_methods.ray"): try: run_LP(args) except Exception: # Expected to fail due to complex initialization pass - + # Verify setup calls mock_check_files.assert_called_once() mock_get_mapping.assert_called_once() @@ -312,14 +318,15 @@ def test_lp_basic_setup(self, mock_server_lp, mock_get_time, @pytest.mark.integration class TestFedGraphEndToEnd: """End-to-end integration tests for complete federated learning workflows.""" - - @patch('fedgraph.federated_methods.data_loader') - @patch('fedgraph.federated_methods.run_NC') + + @patch("fedgraph.federated_methods.data_loader") + @patch("fedgraph.federated_methods.run_NC") def test_complete_nc_pipeline(self, mock_run_nc, mock_data_loader): """Test complete NC pipeline from args to execution.""" - from fedgraph.federated_methods import run_fedgraph import attridict - + + from fedgraph.federated_methods import run_fedgraph + # Create realistic args args = attridict.AttriDict() args.fedgraph_task = "NC" @@ -332,34 +339,38 @@ def test_complete_nc_pipeline(self, mock_run_nc, mock_data_loader): args.use_huggingface = False args.use_lowrank = False args.use_encryption = False - + # Mock data mock_data = ( torch.randn(2, 100), # edge_index - torch.randn(100, 50), # features - torch.randint(0, 7, (100,)), # labels + torch.randn(100, 50), # features + torch.randint(0, 7, (100,)), # labels torch.arange(0, 70), # idx_train - torch.arange(85, 100), # idx_test + torch.arange(85, 100), # idx_test 7, # class_num - [torch.arange(i*30, (i+1)*30) for i in range(3)], # splits - {}, {}, {}, {} # communication data + [torch.arange(i * 30, (i + 1) * 30) for i in range(3)], # splits + {}, + {}, + {}, + {}, # communication data ) mock_data_loader.return_value = mock_data - + # Run the pipeline run_fedgraph(args) - + # Verify calls mock_data_loader.assert_called_once_with(args) mock_run_nc.assert_called_once_with(args, mock_data) - - @patch('fedgraph.federated_methods.data_loader') - @patch('fedgraph.federated_methods.run_GC') + + @patch("fedgraph.federated_methods.data_loader") + @patch("fedgraph.federated_methods.run_GC") def test_complete_gc_pipeline(self, mock_run_gc, mock_data_loader): """Test complete GC pipeline from args to execution.""" - from fedgraph.federated_methods import run_fedgraph import attridict - + + from fedgraph.federated_methods import run_fedgraph + # Create realistic args args = attridict.AttriDict() args.fedgraph_task = "GC" @@ -369,40 +380,57 @@ def test_complete_gc_pipeline(self, mock_run_gc, mock_data_loader): args.num_rounds = 10 args.device = "cpu" args.use_huggingface = False - + # Mock data mock_data = { - "0-PROTEINS": ({"train": Mock(), "val": Mock(), "test": Mock()}, 10, 2, 100), - "1-PROTEINS": ({"train": Mock(), "val": Mock(), "test": Mock()}, 10, 2, 100), - "2-PROTEINS": ({"train": Mock(), "val": Mock(), "test": Mock()}, 10, 2, 100) + "0-PROTEINS": ( + {"train": Mock(), "val": Mock(), "test": Mock()}, + 10, + 2, + 100, + ), + "1-PROTEINS": ( + {"train": Mock(), "val": Mock(), "test": Mock()}, + 10, + 2, + 100, + ), + "2-PROTEINS": ( + {"train": Mock(), "val": Mock(), "test": Mock()}, + 10, + 2, + 100, + ), } mock_data_loader.return_value = mock_data - + # Run the pipeline run_fedgraph(args) - + # Verify calls mock_data_loader.assert_called_once_with(args) mock_run_gc.assert_called_once_with(args, mock_data) - + def test_task_validation_and_routing(self): """Test task validation and proper routing to task-specific functions.""" - from fedgraph.federated_methods import run_fedgraph import attridict - + + from fedgraph.federated_methods import run_fedgraph + # Test invalid task args = attridict.AttriDict() args.fedgraph_task = "INVALID_TASK" - - with patch('fedgraph.federated_methods.data_loader'): + + with patch("fedgraph.federated_methods.data_loader"): with pytest.raises((KeyError, AttributeError)): run_fedgraph(args) - + def test_configuration_validation(self): """Test comprehensive configuration validation.""" - from fedgraph.federated_methods import run_fedgraph import attridict - + + from fedgraph.federated_methods import run_fedgraph + # FedGCN-v2 explicitly supports the low-rank + encryption combination # (encrypted SVD-compressed pretraining), so the configuration # validation now only enforces that low-rank is restricted to NC @@ -412,42 +440,47 @@ def test_configuration_validation(self): args.fedgraph_task = "GC" args.use_lowrank = True - with pytest.raises(ValueError, match="Low-rank compression currently only supported for NC tasks"): + with pytest.raises( + ValueError, + match="Low-rank compression currently only supported for NC tasks", + ): run_fedgraph(args) - + # Test low-rank with wrong task args.method = "FedAvg" args.fedgraph_task = "GC" - - with pytest.raises(ValueError, match="Low-rank compression currently only supported for NC"): + + with pytest.raises( + ValueError, match="Low-rank compression currently only supported for NC" + ): run_fedgraph(args) @pytest.mark.integration class TestComponentInteraction: """Test interaction between different FedGraph components.""" - + def test_data_trainer_server_interaction(self): """Test interaction flow between data processing, trainers, and server.""" # This test verifies that data flows correctly between components - + # 1. Data processing produces correct format edge_index = torch.randint(0, 50, (2, 100)) features = torch.randn(50, 20) labels = torch.randint(0, 3, (50,)) - + # 2. Trainers can initialize with this data - with patch('fedgraph.trainer_class.GCN') as mock_gcn: + with patch("fedgraph.trainer_class.GCN") as mock_gcn: from fedgraph.trainer_class import Trainer_General - + mock_model = Mock() mock_gcn.return_value = mock_model mock_model.to.return_value = mock_model - + trainer = Trainer_General( rank=0, args_hidden=32, - device=torch.device('cpu'), + device=torch.device("cpu"), args=Mock(local_step=3, method="FedAvg"), local_node_index=torch.arange(0, 25), communicate_node_index=torch.arange(0, 40), @@ -456,53 +489,55 @@ def test_data_trainer_server_interaction(self): test_labels=labels[20:30], features=features, idx_train=torch.arange(0, 15), - idx_test=torch.arange(15, 25) + idx_test=torch.arange(15, 25), ) - + # 3. Server can aggregate trainer parameters - with patch('fedgraph.server_class.AggreGCN') as mock_server_gcn: + with patch("fedgraph.server_class.AggreGCN") as mock_server_gcn: from fedgraph.server_class import Server - + mock_server_model = Mock() mock_server_gcn.return_value = mock_server_model mock_server_model.to.return_value = mock_server_model - + server = Server( feature_dim=20, args_hidden=32, class_num=3, - device=torch.device('cpu'), + device=torch.device("cpu"), trainers=[trainer], - args=Mock(num_hops=2, dataset="test", num_layers=2) + args=Mock(num_hops=2, dataset="test", num_layers=2), ) - + # Test parameter flow trainer.update_params = Mock() - mock_server_model.state_dict.return_value = {'weight': torch.randn(32, 20)} - + mock_server_model.state_dict.return_value = { + "weight": torch.randn(32, 20) + } + with patch("fedgraph.server_class.ray.get"): server.broadcast_params(current_global_epoch=1) trainer.update_params.remote.assert_called() - + def test_monitor_integration(self): """Test monitoring system integration.""" from fedgraph.monitor_class import Monitor - + # Create monitor args = Mock() args.monitor = True args.dataset = "test" args.method = "FedAvg" - - with patch('fedgraph.monitor_class.wandb'): + + with patch("fedgraph.monitor_class.wandb"): monitor = Monitor(args) - + # Test logging functionality monitor.log_metrics = Mock() test_metrics = {"accuracy": 0.85, "loss": 0.3} - + # This would be called during training - if hasattr(monitor, 'log_metrics'): + if hasattr(monitor, "log_metrics"): monitor.log_metrics(test_metrics) monitor.log_metrics.assert_called_with(test_metrics) @@ -511,55 +546,55 @@ def test_monitor_integration(self): @pytest.mark.slow class TestLargeScaleIntegration: """Integration tests for larger scale scenarios.""" - + def test_multi_trainer_coordination(self): """Test coordination between multiple trainers.""" num_trainers = 5 - + # Create multiple trainers trainers = [] for i in range(num_trainers): - with patch('fedgraph.trainer_class.GCN'): + with patch("fedgraph.trainer_class.GCN"): from fedgraph.trainer_class import Trainer_General - + trainer = Trainer_General( rank=i, args_hidden=32, - device=torch.device('cpu'), + device=torch.device("cpu"), args=Mock(local_step=3, method="FedAvg"), - local_node_index=torch.arange(i*10, (i+1)*10), + local_node_index=torch.arange(i * 10, (i + 1) * 10), communicate_node_index=torch.arange(0, 50), adj=torch.randint(0, 50, (2, 50)), train_labels=torch.randint(0, 3, (10,)), test_labels=torch.randint(0, 3, (5,)), features=torch.randn(50, 20), idx_train=torch.arange(0, 8), - idx_test=torch.arange(8, 10) + idx_test=torch.arange(8, 10), ) trainers.append(trainer) - + # Test that all trainers have unique ranks ranks = [trainer.get_rank() for trainer in trainers] assert len(set(ranks)) == num_trainers assert sorted(ranks) == list(range(num_trainers)) - + def test_memory_efficiency(self): """Test memory efficiency with larger data structures.""" # Create larger tensors to test memory handling large_features = torch.randn(1000, 100) large_edge_index = torch.randint(0, 1000, (2, 5000)) - + # Test that operations complete without memory errors from fedgraph.utils_nc import intersect1d, setdiff1d - + t1 = torch.arange(0, 500) t2 = torch.arange(250, 750) - + intersection = intersect1d(t1, t2) difference = setdiff1d(t1, t2) - + assert len(intersection) > 0 assert len(difference) > 0 - + # Clean up del large_features, large_edge_index, intersection, difference diff --git a/tests/test_smoke_e2e.py b/tests/test_smoke_e2e.py index 1974033..2ef6240 100644 --- a/tests/test_smoke_e2e.py +++ b/tests/test_smoke_e2e.py @@ -19,6 +19,7 @@ test accuracy. We do not regression-test exact numbers because the gcn_v2 work intentionally changes the default normalization. """ + from __future__ import annotations import contextlib @@ -31,11 +32,11 @@ from tests.conftest import needs_openfhe, needs_tenseal - # --------------------------------------------------------------------------- # Shared minimal config # --------------------------------------------------------------------------- + def _base_cora_config(**overrides): cfg = { "fedgraph_task": "NC", @@ -67,7 +68,7 @@ def _base_cora_config(**overrides): # try to grab 20 GB of /dev/shm. "ray_init_kwargs": { "num_cpus": 2, - "object_store_memory": 128 * 1024 ** 2, # 128 MiB + "object_store_memory": 128 * 1024**2, # 128 MiB "include_dashboard": False, "configure_logging": False, }, @@ -91,10 +92,11 @@ def _run(cfg) -> float: captured stdout. We tolerate a wide range -- the goal is to verify the pipeline doesn't crash and produces a non-trivial number, not to benchmark.""" - from fedgraph.federated_methods import run_fedgraph import io from contextlib import redirect_stdout + from fedgraph.federated_methods import run_fedgraph + buf = io.StringIO() with redirect_stdout(buf): run_fedgraph(cfg) @@ -114,6 +116,7 @@ def _run(cfg) -> float: # Plaintext path -- original aggregation # --------------------------------------------------------------------------- + @pytest.mark.timeout(120) def test_smoke_plaintext_default_norm_none(): with _ray_shutdown_after(): @@ -126,6 +129,7 @@ def test_smoke_plaintext_default_norm_none(): # Plaintext path -- new (FedGCN-v2) symmetric normalization opt-in # --------------------------------------------------------------------------- + @pytest.mark.timeout(120) def test_smoke_plaintext_norm_sym_opt_in(): with _ray_shutdown_after(): @@ -137,6 +141,7 @@ def test_smoke_plaintext_norm_sym_opt_in(): # TenSEAL backend # --------------------------------------------------------------------------- + @needs_tenseal @pytest.mark.timeout(180) def test_smoke_tenseal_encrypted(): @@ -150,6 +155,7 @@ def test_smoke_tenseal_encrypted(): # OpenFHE threshold backend # --------------------------------------------------------------------------- + @needs_openfhe @pytest.mark.timeout(300) def test_smoke_openfhe_threshold_encrypted(): @@ -166,6 +172,7 @@ def test_smoke_openfhe_threshold_encrypted(): # OpenFHE threshold + low-rank # --------------------------------------------------------------------------- + @needs_openfhe @pytest.mark.timeout(300) def test_smoke_openfhe_threshold_lowrank(): diff --git a/tests/test_smoke_unit.py b/tests/test_smoke_unit.py index dbc467e..c2bbd97 100644 --- a/tests/test_smoke_unit.py +++ b/tests/test_smoke_unit.py @@ -5,6 +5,7 @@ a laptop CPU; the end-to-end pipeline is covered in ``test_smoke_e2e.py``. """ + from __future__ import annotations import pytest @@ -12,11 +13,11 @@ from fedgraph.utils_nc import get_1hop_feature_sum - # --------------------------------------------------------------------------- # Adjacency normalization (backward-compat default + new opt-in modes) # --------------------------------------------------------------------------- + def _toy_graph(): # 4-node line: 0-1-2-3 edge_index = torch.tensor( @@ -39,23 +40,16 @@ def test_default_norm_type_is_none_backward_compat(): def test_norm_type_sym_matches_gcn_normalization(): features, edge_index = _toy_graph() - out = get_1hop_feature_sum( - features, edge_index, device="cpu", norm_type="sym" - ) + out = get_1hop_feature_sum(features, edge_index, device="cpu", norm_type="sym") # Row 0 has degree 2 (self + 1 neighbour); neighbour 1 has degree 3 # (self + 0 + 2). Symmetric weight = 1/sqrt(deg_i * deg_j). - expected = ( - features[0] / 2.0 - + features[1] / (2 ** 0.5 * 3 ** 0.5) - ) + expected = features[0] / 2.0 + features[1] / (2**0.5 * 3**0.5) assert torch.allclose(out[0], expected, atol=1e-6) def test_norm_type_row_is_stochastic(): features, edge_index = _toy_graph() - out = get_1hop_feature_sum( - features, edge_index, device="cpu", norm_type="row" - ) + out = get_1hop_feature_sum(features, edge_index, device="cpu", norm_type="row") # Each output row is a convex combination of its self-loop neighbourhood. row_sums = out.sum(dim=1) assert torch.allclose(row_sums, torch.ones_like(row_sums), atol=1e-6) @@ -64,15 +58,14 @@ def test_norm_type_row_is_stochastic(): def test_unknown_norm_type_raises(): features, edge_index = _toy_graph() with pytest.raises(ValueError): - get_1hop_feature_sum( - features, edge_index, device="cpu", norm_type="bogus" - ) + get_1hop_feature_sum(features, edge_index, device="cpu", norm_type="bogus") # --------------------------------------------------------------------------- # Low-rank compression (round-trip) # --------------------------------------------------------------------------- + def test_svd_round_trip_is_close_for_low_rank_matrix(): from fedgraph.low_rank.compression_utils import svd_compress, svd_decompress @@ -102,6 +95,7 @@ def test_svd_handles_rank_larger_than_min_shape(): # Optional OpenFHE wrapper smoke test # --------------------------------------------------------------------------- + def test_openfhe_wrapper_imports_or_skips(): """The wrapper module must not raise at import time even when openfhe is missing on the host.""" diff --git a/tests/test_threshold_ckks_min.py b/tests/test_threshold_ckks_min.py index 4995b7f..2b6dc18 100644 --- a/tests/test_threshold_ckks_min.py +++ b/tests/test_threshold_ckks_min.py @@ -1,7 +1,9 @@ # tests/test_threshold_ckks_min.py -import openfhe import math +import openfhe + + def make_cc(): params = openfhe.CCParamsCKKSRNS() params.SetSecurityLevel(openfhe.HEStd_128_classic) @@ -21,6 +23,7 @@ def make_cc(): cc.Enable(feature) return cc + def test_two_party_threshold_ckks_add(): cc = make_cc() @@ -50,12 +53,13 @@ def test_two_party_threshold_ckks_add(): fused = cc.MultipartyDecryptFusion([p_lead, p_main]) out = fused.GetRealPackedValue() - expect = [a+b for a,b in zip(x,y)] + expect = [a + b for a, b in zip(x, y)] print(f"Expected: {expect}") print(f"Result: {out[:len(expect)]}") - assert all(abs(e-r) < 1e-3 for e,r in zip(expect, out[:len(expect)])) + assert all(abs(e - r) < 1e-3 for e, r in zip(expect, out[: len(expect)])) print("✅ Two-party threshold CKKS test passed!") + if __name__ == "__main__": test_two_party_threshold_ckks_add() diff --git a/tests/unit/test_data_process.py b/tests/unit/test_data_process.py index 802ea08..ec538eb 100644 --- a/tests/unit/test_data_process.py +++ b/tests/unit/test_data_process.py @@ -1,92 +1,93 @@ -import pytest -import torch -import numpy as np -import tempfile import os import pickle as pkl -from unittest.mock import Mock, patch, MagicMock, mock_open +import tempfile +from unittest.mock import MagicMock, Mock, mock_open, patch + import attridict +import numpy as np +import pytest +import torch import torch_sparse from torch_geometric.data import Data from torch_geometric.datasets import TUDataset from torch_geometric.loader import DataLoader from fedgraph.data_process import ( + GC_rand_split_chunk, + NC_load_data, + NC_parse_index_file, data_loader, - data_loader_NC, data_loader_GC, - NC_parse_index_file, - download_file_from_github, - NC_load_data, - GC_rand_split_chunk, + data_loader_GC_multiple, data_loader_GC_single, - data_loader_GC_multiple + data_loader_NC, + download_file_from_github, ) class TestDataLoader: """Test the main data_loader function.""" - + def test_data_loader_LP_returns_none(self): """Test that data_loader returns None for LP task.""" args = Mock() args.fedgraph_task = "LP" - + result = data_loader(args) assert result is None - - @patch('fedgraph.data_process.data_loader_NC') + + @patch("fedgraph.data_process.data_loader_NC") def test_data_loader_NC_calls_correct_function(self, mock_nc_loader): """Test that data_loader calls data_loader_NC for NC task.""" args = Mock() args.fedgraph_task = "NC" mock_nc_loader.return_value = "mock_nc_data" - + result = data_loader(args) - + mock_nc_loader.assert_called_once_with(args) assert result == "mock_nc_data" - - @patch('fedgraph.data_process.data_loader_GC') + + @patch("fedgraph.data_process.data_loader_GC") def test_data_loader_GC_calls_correct_function(self, mock_gc_loader): """Test that data_loader calls data_loader_GC for GC task.""" args = Mock() args.fedgraph_task = "GC" mock_gc_loader.return_value = "mock_gc_data" - + result = data_loader(args) - + mock_gc_loader.assert_called_once_with(args) assert result == "mock_gc_data" class TestNCParseIndexFile: """Test NC_parse_index_file function.""" - + def test_parse_index_file_success(self, temp_dir): """Test parsing index file successfully.""" # Create test file test_file = os.path.join(temp_dir, "test_index.txt") - with open(test_file, 'w') as f: + with open(test_file, "w") as f: f.write("1\n2\n3\n10\n") - + result = NC_parse_index_file(test_file) assert result == [1, 2, 3, 10] - + def test_parse_empty_file(self, temp_dir): """Test parsing empty index file.""" test_file = os.path.join(temp_dir, "empty.txt") - with open(test_file, 'w') as f: + with open(test_file, "w") as f: pass - + result = NC_parse_index_file(test_file) assert result == [] class TestDownloadFileFromGithub: """Test download_file_from_github function.""" - - @patch('fedgraph.data_process.requests.get') + + @patch("fedgraph.data_process.requests.get") def test_download_file_success(self, mock_get, temp_dir): """Test successful file download.""" save_path = os.path.join(temp_dir, "test_file.txt") @@ -94,45 +95,46 @@ def test_download_file_success(self, mock_get, temp_dir): mock_response.status_code = 200 mock_response.iter_content.return_value = [b"test", b"data"] mock_get.return_value = mock_response - + download_file_from_github("http://test.com/file.txt", save_path) - + assert os.path.exists(save_path) - with open(save_path, 'rb') as f: + with open(save_path, "rb") as f: assert f.read() == b"testdata" - - @patch('fedgraph.data_process.requests.get') + + @patch("fedgraph.data_process.requests.get") def test_download_file_failure(self, mock_get, temp_dir): """Test failed file download.""" save_path = os.path.join(temp_dir, "test_file.txt") mock_response = Mock() mock_response.status_code = 404 mock_get.return_value = mock_response - + with pytest.raises(Exception, match="Failed to download"): download_file_from_github("http://test.com/nonexistent.txt", save_path) - + def test_file_already_exists(self, temp_dir): """Test that existing files are not re-downloaded.""" save_path = os.path.join(temp_dir, "existing_file.txt") - with open(save_path, 'w') as f: + with open(save_path, "w") as f: f.write("existing content") - - with patch('fedgraph.data_process.requests.get') as mock_get: + + with patch("fedgraph.data_process.requests.get") as mock_get: download_file_from_github("http://test.com/file.txt", save_path) mock_get.assert_not_called() class TestNCLoadData: """Test NC_load_data function.""" - - @patch('fedgraph.data_process.download_file_from_github') - @patch('builtins.open', new_callable=mock_open) - @patch('fedgraph.data_process.pkl.load') - @patch('fedgraph.data_process.NC_parse_index_file') - @patch('os.makedirs') - def test_load_cora_dataset(self, mock_makedirs, mock_parse, mock_pkl_load, - mock_open_file, mock_download): + + @patch("fedgraph.data_process.download_file_from_github") + @patch("builtins.open", new_callable=mock_open) + @patch("fedgraph.data_process.pkl.load") + @patch("fedgraph.data_process.NC_parse_index_file") + @patch("os.makedirs") + def test_load_cora_dataset( + self, mock_makedirs, mock_parse, mock_pkl_load, mock_open_file, mock_download + ): """Test loading Cora dataset.""" # Setup mock data mock_x = np.random.random((100, 50)) @@ -141,29 +143,33 @@ def test_load_cora_dataset(self, mock_makedirs, mock_parse, mock_pkl_load, mock_ty = np.random.randint(0, 7, (1000, 7)) mock_allx = np.random.random((1708, 50)) mock_ally = np.random.randint(0, 7, (1708, 7)) - mock_graph = {i: [i+1] for i in range(1707)} - + mock_graph = {i: [i + 1] for i in range(1707)} + mock_pkl_load.side_effect = [ - mock_x, mock_y, mock_tx, mock_ty, - mock_allx, mock_ally, mock_graph + mock_x, + mock_y, + mock_tx, + mock_ty, + mock_allx, + mock_ally, + mock_graph, ] mock_parse.return_value = list(range(1000, 1708)) - + # Mock scipy sparse matrices - with patch('scipy.sparse.vstack') as mock_vstack, \ - patch('networkx.adjacency_matrix') as mock_adj, \ - patch('networkx.from_dict_of_lists') as mock_from_dict: - + with patch("scipy.sparse.vstack") as mock_vstack, patch( + "networkx.adjacency_matrix" + ) as mock_adj, patch("networkx.from_dict_of_lists") as mock_from_dict: mock_features = Mock() mock_features.toarray.return_value = np.random.random((2708, 50)) mock_vstack.return_value = mock_features - + mock_adj_matrix = Mock() mock_adj_matrix.toarray.return_value = np.random.random((2708, 2708)) mock_adj.return_value = mock_adj_matrix - + result = NC_load_data("cora") - + assert len(result) == 6 features, adj, labels, idx_train, idx_val, idx_test = result assert isinstance(features, torch.Tensor) @@ -176,29 +182,29 @@ def test_load_cora_dataset(self, mock_makedirs, mock_parse, mock_pkl_load, class TestGCRandSplitChunk: """Test GC_rand_split_chunk function.""" - + def test_non_overlapping_split(self): """Test non-overlapping split of graphs.""" graphs = [Mock() for _ in range(100)] num_trainer = 5 - + result = GC_rand_split_chunk(graphs, num_trainer, overlap=False, seed=42) - + assert len(result) == num_trainer # Check that all chunks together contain all graphs total_graphs = sum(len(chunk) for chunk in result) assert total_graphs == len(graphs) - + def test_overlapping_split(self): """Test overlapping split of graphs.""" graphs = [Mock() for _ in range(100)] num_trainer = 3 - - with patch('numpy.random.randint') as mock_randint: + + with patch("numpy.random.randint") as mock_randint: mock_randint.return_value = np.array([75, 80, 85]) - + result = GC_rand_split_chunk(graphs, num_trainer, overlap=True, seed=42) - + assert len(result) == num_trainer for chunk in result: assert len(chunk) >= 50 # minimum size from function @@ -206,14 +212,15 @@ def test_overlapping_split(self): class TestDataLoaderGCSingle: """Test data_loader_GC_single function.""" - - @patch('fedgraph.data_process.TUDataset') - @patch('fedgraph.data_process.GC_rand_split_chunk') - @patch('fedgraph.data_process.split_data') - @patch('fedgraph.data_process.DataLoader') - @patch('fedgraph.data_process.get_num_graph_labels') - def test_data_loader_GC_single_success(self, mock_get_labels, mock_dataloader, - mock_split, mock_chunk, mock_tudataset): + + @patch("fedgraph.data_process.TUDataset") + @patch("fedgraph.data_process.GC_rand_split_chunk") + @patch("fedgraph.data_process.split_data") + @patch("fedgraph.data_process.DataLoader") + @patch("fedgraph.data_process.get_num_graph_labels") + def test_data_loader_GC_single_success( + self, mock_get_labels, mock_dataloader, mock_split, mock_chunk, mock_tudataset + ): """Test successful GC single dataset loading.""" # Setup mocks mock_dataset = Mock() @@ -221,23 +228,19 @@ def test_data_loader_GC_single_success(self, mock_get_labels, mock_dataloader, mock_graph.num_node_features = 10 mock_dataset.__iter__ = Mock(return_value=iter([mock_graph] * 100)) mock_tudataset.return_value = mock_dataset - + mock_chunk.return_value = [[mock_graph] * 20 for _ in range(5)] mock_split.side_effect = [ ([mock_graph] * 16, [mock_graph] * 4), # train/val_test split - ([mock_graph] * 2, [mock_graph] * 2) # val/test split + ([mock_graph] * 2, [mock_graph] * 2), # val/test split ] mock_get_labels.return_value = 2 mock_dataloader.return_value = Mock() - + result = data_loader_GC_single( - datapath="/tmp", - dataset="PROTEINS", - num_trainer=5, - batch_size=32, - seed=42 + datapath="/tmp", dataset="PROTEINS", num_trainer=5, batch_size=32, seed=42 ) - + assert isinstance(result, dict) assert len(result) == 5 # num_trainer for key, value in result.items(): @@ -247,13 +250,14 @@ def test_data_loader_GC_single_success(self, mock_get_labels, mock_dataloader, class TestDataLoaderGCMultiple: """Test data_loader_GC_multiple function.""" - - @patch('fedgraph.data_process.TUDataset') - @patch('fedgraph.data_process.split_data') - @patch('fedgraph.data_process.DataLoader') - @patch('fedgraph.data_process.get_num_graph_labels') - def test_data_loader_GC_multiple_small_group(self, mock_get_labels, mock_dataloader, - mock_split, mock_tudataset): + + @patch("fedgraph.data_process.TUDataset") + @patch("fedgraph.data_process.split_data") + @patch("fedgraph.data_process.DataLoader") + @patch("fedgraph.data_process.get_num_graph_labels") + def test_data_loader_GC_multiple_small_group( + self, mock_get_labels, mock_dataloader, mock_split, mock_tudataset + ): """Test GC multiple datasets loading for small group.""" # Setup mocks mock_dataset = Mock() @@ -261,26 +265,32 @@ def test_data_loader_GC_multiple_small_group(self, mock_get_labels, mock_dataloa mock_graph.num_node_features = 10 mock_dataset.__iter__ = Mock(return_value=iter([mock_graph] * 50)) mock_tudataset.return_value = mock_dataset - + mock_split.side_effect = [ ([mock_graph] * 40, [mock_graph] * 10), # train/val_test split - ([mock_graph] * 5, [mock_graph] * 5) # val/test split + ([mock_graph] * 5, [mock_graph] * 5), # val/test split ] * 8 # For 8 datasets in small group - + mock_get_labels.return_value = 2 mock_dataloader.return_value = Mock() - + result = data_loader_GC_multiple( - datapath="/tmp", - dataset_group="small", - batch_size=32, - seed=42 + datapath="/tmp", dataset_group="small", batch_size=32, seed=42 ) - + assert isinstance(result, dict) - expected_datasets = ["MUTAG", "BZR", "COX2", "DHFR", "PTC_MR", "ENZYMES", "DD", "PROTEINS"] + expected_datasets = [ + "MUTAG", + "BZR", + "COX2", + "DHFR", + "PTC_MR", + "ENZYMES", + "DD", + "PROTEINS", + ] assert len(result) == len(expected_datasets) - + for dataset_name in expected_datasets: assert dataset_name in result value = result[dataset_name] @@ -290,12 +300,13 @@ def test_data_loader_GC_multiple_small_group(self, mock_get_labels, mock_dataloa class TestDataLoaderNC: """Test data_loader_NC function with integration.""" - - @patch('fedgraph.data_process.NC_load_data') - @patch('fedgraph.data_process.label_dirichlet_partition') - @patch('fedgraph.data_process.get_in_comm_indexes') - def test_data_loader_NC_complete_flow(self, mock_get_indexes, mock_partition, - mock_load_data): + + @patch("fedgraph.data_process.NC_load_data") + @patch("fedgraph.data_process.label_dirichlet_partition") + @patch("fedgraph.data_process.get_in_comm_indexes") + def test_data_loader_NC_complete_flow( + self, mock_get_indexes, mock_partition, mock_load_data + ): """Test complete NC data loading flow.""" # Setup mock args args = Mock() @@ -305,12 +316,12 @@ def test_data_loader_NC_complete_flow(self, mock_get_indexes, mock_partition, args.iid_beta = 0.5 args.distribution_type = "dirichlet" args.num_hops = 2 - + # Setup mock data num_nodes = 100 num_features = 50 num_classes = 7 - + features = torch.randn(num_nodes, num_features) adj = torch_sparse.tensor.SparseTensor.from_dense( torch.randn(num_nodes, num_nodes) @@ -319,32 +330,55 @@ def test_data_loader_NC_complete_flow(self, mock_get_indexes, mock_partition, idx_train = torch.arange(0, 70) idx_val = torch.arange(70, 85) idx_test = torch.arange(85, 100) - - mock_load_data.return_value = (features, adj, labels, idx_train, idx_val, idx_test) - + + mock_load_data.return_value = ( + features, + adj, + labels, + idx_train, + idx_val, + idx_test, + ) + # Mock partition results mock_partition.return_value = [ list(range(0, 33)), list(range(33, 66)), - list(range(66, 100)) + list(range(66, 100)), ] - + # Mock communication indexes mock_get_indexes.return_value = ( - {i: torch.arange(i*10, (i+1)*10) for i in range(3)}, # communicate_node_global_indexes - {i: torch.arange(0, 5) for i in range(3)}, # in_com_train_node_local_indexes - {i: torch.arange(5, 8) for i in range(3)}, # in_com_test_node_local_indexes - {i: torch.stack([torch.arange(0, 10), torch.arange(1, 11)]) for i in range(3)} # global_edge_indexes_clients + { + i: torch.arange(i * 10, (i + 1) * 10) for i in range(3) + }, # communicate_node_global_indexes + { + i: torch.arange(0, 5) for i in range(3) + }, # in_com_train_node_local_indexes + {i: torch.arange(5, 8) for i in range(3)}, # in_com_test_node_local_indexes + { + i: torch.stack([torch.arange(0, 10), torch.arange(1, 11)]) + for i in range(3) + }, # global_edge_indexes_clients ) - + result = data_loader_NC(args) - + assert len(result) == 11 - (edge_index, returned_features, returned_labels, returned_idx_train, - returned_idx_test, class_num, split_node_indexes, - communicate_node_global_indexes, in_com_train_node_local_indexes, - in_com_test_node_local_indexes, global_edge_indexes_clients) = result - + ( + edge_index, + returned_features, + returned_labels, + returned_idx_train, + returned_idx_test, + class_num, + split_node_indexes, + communicate_node_global_indexes, + in_com_train_node_local_indexes, + in_com_test_node_local_indexes, + global_edge_indexes_clients, + ) = result + # Verify results assert torch.equal(returned_features, features) assert torch.equal(returned_labels, labels) @@ -352,7 +386,7 @@ def test_data_loader_NC_complete_flow(self, mock_get_indexes, mock_partition, assert torch.equal(returned_idx_test, idx_test) assert class_num == num_classes assert len(split_node_indexes) == args.n_trainer - + # Verify function calls mock_load_data.assert_called_once_with(args.dataset) mock_partition.assert_called_once() @@ -361,8 +395,8 @@ def test_data_loader_NC_complete_flow(self, mock_get_indexes, mock_partition, class TestDataLoaderGC: """Test data_loader_GC function.""" - - @patch('fedgraph.data_process.data_loader_GC_multiple') + + @patch("fedgraph.data_process.data_loader_GC_multiple") def test_data_loader_GC_multiple_datasets(self, mock_multiple): """Test GC data loader for multiple datasets.""" args = Mock() @@ -372,21 +406,21 @@ def test_data_loader_GC_multiple_datasets(self, mock_multiple): args.batch_size = 32 args.convert_x = False args.seed_split_data = 42 - + mock_multiple.return_value = {"test": "data"} - + result = data_loader_GC(args) - + mock_multiple.assert_called_once_with( datapath=args.datapath, dataset_group=args.dataset_group, batch_size=args.batch_size, convert_x=args.convert_x, - seed=args.seed_split_data + seed=args.seed_split_data, ) assert result == {"test": "data"} - - @patch('fedgraph.data_process.data_loader_GC_single') + + @patch("fedgraph.data_process.data_loader_GC_single") def test_data_loader_GC_single_dataset(self, mock_single): """Test GC data loader for single dataset.""" args = Mock() @@ -398,11 +432,11 @@ def test_data_loader_GC_single_dataset(self, mock_single): args.convert_x = False args.seed_split_data = 42 args.overlap = False - + mock_single.return_value = {"test": "data"} - + result = data_loader_GC(args) - + mock_single.assert_called_once_with( datapath=args.datapath, dataset=args.dataset, @@ -410,7 +444,7 @@ def test_data_loader_GC_single_dataset(self, mock_single): batch_size=args.batch_size, convert_x=args.convert_x, seed=args.seed_split_data, - overlap=args.overlap + overlap=args.overlap, ) assert result == {"test": "data"} @@ -418,13 +452,14 @@ def test_data_loader_GC_single_dataset(self, mock_single): @pytest.mark.integration class TestDataProcessIntegration: """Integration tests for data processing functions.""" - + def test_complete_NC_workflow_mock(self): """Test complete NC workflow with mocked dependencies.""" - with patch('fedgraph.data_process.NC_load_data') as mock_load, \ - patch('fedgraph.data_process.label_dirichlet_partition') as mock_partition, \ - patch('fedgraph.data_process.get_in_comm_indexes') as mock_indexes: - + with patch("fedgraph.data_process.NC_load_data") as mock_load, patch( + "fedgraph.data_process.label_dirichlet_partition" + ) as mock_partition, patch( + "fedgraph.data_process.get_in_comm_indexes" + ) as mock_indexes: # Setup test data features = torch.randn(100, 50) adj = torch_sparse.tensor.SparseTensor.from_dense(torch.eye(100)) @@ -432,11 +467,18 @@ def test_complete_NC_workflow_mock(self): idx_train = torch.arange(0, 70) idx_val = torch.arange(70, 85) idx_test = torch.arange(85, 100) - - mock_load.return_value = (features, adj, labels, idx_train, idx_val, idx_test) + + mock_load.return_value = ( + features, + adj, + labels, + idx_train, + idx_val, + idx_test, + ) mock_partition.return_value = [list(range(0, 50)), list(range(50, 100))] mock_indexes.return_value = ({}, {}, {}, {}) - + # Create args args = Mock() args.fedgraph_task = "NC" @@ -446,12 +488,12 @@ def test_complete_NC_workflow_mock(self): args.iid_beta = 0.5 args.distribution_type = "dirichlet" args.num_hops = 2 - + # Test the complete workflow result = data_loader(args) - + assert result is not None assert len(result) == 11 mock_load.assert_called_once() mock_partition.assert_called_once() - mock_indexes.assert_called_once() \ No newline at end of file + mock_indexes.assert_called_once() diff --git a/tests/unit/test_federated_methods.py b/tests/unit/test_federated_methods.py index acb1461..c8239f7 100644 --- a/tests/unit/test_federated_methods.py +++ b/tests/unit/test_federated_methods.py @@ -1,17 +1,18 @@ +from unittest.mock import MagicMock, Mock, patch + +import attridict +import numpy as np import pytest import torch -import numpy as np -from unittest.mock import Mock, patch, MagicMock -import attridict from fedgraph.federated_methods import ( _resolve_nc_class_num, _resolve_nc_global_node_num, run_fedgraph, run_fedgraph_enhanced, - run_NC, run_GC, - run_LP + run_LP, + run_NC, ) @@ -47,7 +48,9 @@ def test_rejects_inconsistent_huggingface_class_num_metadata(self): def test_rejects_huggingface_data_without_labels(self): trainer_information = [{"label_num": None}, {"label_num": None}] - with pytest.raises(ValueError, match="all train and test label tensors are empty"): + with pytest.raises( + ValueError, match="all train and test label tensors are empty" + ): _resolve_nc_class_num(True, trainer_information) @@ -82,7 +85,7 @@ def test_rejects_inconsistent_huggingface_global_node_num_metadata(self): class TestRunFedgraph: """Test run_fedgraph main orchestration function.""" - + def setup_method(self): """Setup test data for NC tests.""" self.args = attridict.AttriDict() @@ -92,66 +95,71 @@ def setup_method(self): self.args.use_encryption = False self.args.use_huggingface = False self.args.num_hops = 2 - - @patch('fedgraph.federated_methods.data_loader') - @patch('fedgraph.federated_methods.run_NC') + + @patch("fedgraph.federated_methods.data_loader") + @patch("fedgraph.federated_methods.run_NC") def test_run_fedgraph_nc_task(self, mock_run_nc, mock_data_loader): """Test run_fedgraph with NC task.""" mock_data = Mock() mock_data_loader.return_value = mock_data - + run_fedgraph(self.args) - + mock_data_loader.assert_called_once_with(self.args) mock_run_nc.assert_called_once_with(self.args, mock_data) - - @patch('fedgraph.federated_methods.data_loader') - @patch('fedgraph.federated_methods.run_GC') + + @patch("fedgraph.federated_methods.data_loader") + @patch("fedgraph.federated_methods.run_GC") def test_run_fedgraph_gc_task(self, mock_run_gc, mock_data_loader): """Test run_fedgraph with GC task.""" self.args.fedgraph_task = "GC" mock_data = Mock() mock_data_loader.return_value = mock_data - + run_fedgraph(self.args) - + mock_data_loader.assert_called_once_with(self.args) mock_run_gc.assert_called_once_with(self.args, mock_data) - - @patch('fedgraph.federated_methods.run_LP') + + @patch("fedgraph.federated_methods.run_LP") def test_run_fedgraph_lp_task(self, mock_run_lp): """Test run_fedgraph with LP task.""" self.args.fedgraph_task = "LP" - + run_fedgraph(self.args) - + mock_run_lp.assert_called_once_with(self.args) - - @patch('fedgraph.federated_methods.run_NC_lowrank') + + @patch("fedgraph.federated_methods.run_NC_lowrank") def test_run_fedgraph_nc_lowrank(self, mock_run_nc_lowrank): """Test run_fedgraph with NC task and low-rank compression.""" self.args.use_lowrank = True - - with patch('fedgraph.federated_methods.data_loader') as mock_data_loader: + + with patch("fedgraph.federated_methods.data_loader") as mock_data_loader: mock_data = Mock() mock_data_loader.return_value = mock_data - + run_fedgraph(self.args) - + mock_run_nc_lowrank.assert_called_once_with(self.args, mock_data) - + def test_run_fedgraph_lowrank_validation_nc_only(self): """Test that low-rank compression only works with NC tasks.""" self.args.use_lowrank = True self.args.fedgraph_task = "GC" - - with pytest.raises(ValueError, match="Low-rank compression currently only supported for NC tasks"): + + with pytest.raises( + ValueError, + match="Low-rank compression currently only supported for NC tasks", + ): run_fedgraph(self.args) - - @patch('fedgraph.federated_methods.data_loader') - @patch('fedgraph.federated_methods.run_NC_lowrank') + + @patch("fedgraph.federated_methods.data_loader") + @patch("fedgraph.federated_methods.run_NC_lowrank") def test_run_fedgraph_lowrank_works_with_non_fedavg_method( - self, mock_run_nc_lowrank, mock_data_loader, + self, + mock_run_nc_lowrank, + mock_data_loader, ): """Low-rank compression is no longer restricted to FedAvg. @@ -166,11 +174,13 @@ def test_run_fedgraph_lowrank_works_with_non_fedavg_method( run_fedgraph(self.args) mock_run_nc_lowrank.assert_called_once() - - @patch('fedgraph.federated_methods.data_loader') - @patch('fedgraph.federated_methods.run_NC') + + @patch("fedgraph.federated_methods.data_loader") + @patch("fedgraph.federated_methods.run_NC") def test_run_fedgraph_lowrank_with_openfhe_dispatches_to_run_nc( - self, mock_run_nc, mock_data_loader, + self, + mock_run_nc, + mock_data_loader, ): """Combining low-rank with the OpenFHE threshold backend is the FedGCN-v2 path. It dispatches to run_NC (which carries the @@ -183,7 +193,7 @@ def test_run_fedgraph_lowrank_with_openfhe_dispatches_to_run_nc( run_fedgraph(self.args) mock_run_nc.assert_called_once() - @patch('fedgraph.federated_methods.data_loader') + @patch("fedgraph.federated_methods.data_loader") def test_run_fedgraph_rejects_unsupported_nc_num_hops(self, mock_data_loader): """Test that ambiguous 1-hop NC mode is rejected before data loading.""" self.args.num_hops = 1 @@ -192,36 +202,36 @@ def test_run_fedgraph_rejects_unsupported_nc_num_hops(self, mock_data_loader): run_fedgraph(self.args) mock_data_loader.assert_not_called() - - @patch('fedgraph.federated_methods.data_loader') - @patch('fedgraph.federated_methods.run_NC') + + @patch("fedgraph.federated_methods.data_loader") + @patch("fedgraph.federated_methods.run_NC") def test_run_fedgraph_huggingface_nc(self, mock_run_nc, mock_data_loader): """Test run_fedgraph with NC task using Hugging Face.""" self.args.use_huggingface = True - + run_fedgraph(self.args) - + mock_data_loader.assert_not_called() # Should not load data when using HF mock_run_nc.assert_called_once_with(self.args, None) class TestRunFedgraphEnhanced: """Test run_fedgraph_enhanced function.""" - - @patch('fedgraph.federated_methods.run_fedgraph') + + @patch("fedgraph.federated_methods.run_fedgraph") def test_run_fedgraph_enhanced_calls_main(self, mock_run_fedgraph): """Test that run_fedgraph_enhanced calls the main function.""" args = attridict.AttriDict() args.fedgraph_task = "NC" - + run_fedgraph_enhanced(args) - + mock_run_fedgraph.assert_called_once_with(args) class TestRunNC: """Test run_NC function for node classification.""" - + def setup_method(self): """Setup test data for each test method.""" self.args = attridict.AttriDict() @@ -236,85 +246,90 @@ def setup_method(self): self.args.use_ray = True self.args.he = False self.args.dp = False - + # Mock data self.mock_data = ( torch.randn(2, 100), # edge_index - torch.randn(100, 50), # features - torch.randint(0, 7, (100,)), # labels + torch.randn(100, 50), # features + torch.randint(0, 7, (100,)), # labels torch.arange(0, 70), # idx_train - torch.arange(85, 100), # idx_test + torch.arange(85, 100), # idx_test 7, # class_num - [torch.arange(i*10, (i+1)*10) for i in range(3)], # split_node_indexes - {i: torch.arange(i*20, (i+1)*20) for i in range(3)}, # communicate_node_global_indexes - {i: torch.arange(0, 5) for i in range(3)}, # in_com_train_node_local_indexes - {i: torch.arange(5, 8) for i in range(3)}, # in_com_test_node_local_indexes - {i: torch.randn(2, 20) for i in range(3)} # global_edge_indexes_clients + [ + torch.arange(i * 10, (i + 1) * 10) for i in range(3) + ], # split_node_indexes + { + i: torch.arange(i * 20, (i + 1) * 20) for i in range(3) + }, # communicate_node_global_indexes + { + i: torch.arange(0, 5) for i in range(3) + }, # in_com_train_node_local_indexes + {i: torch.arange(5, 8) for i in range(3)}, # in_com_test_node_local_indexes + {i: torch.randn(2, 20) for i in range(3)}, # global_edge_indexes_clients ) - - @patch('fedgraph.federated_methods.ray') - @patch('fedgraph.federated_methods.Monitor') + + @patch("fedgraph.federated_methods.ray") + @patch("fedgraph.federated_methods.Monitor") def test_run_nc_basic_setup(self, mock_monitor, mock_ray): """Test basic setup of run_NC function.""" mock_monitor_instance = Mock() mock_monitor.return_value = mock_monitor_instance - + # Mock Ray mock_ray.init = Mock() mock_ray.get = Mock() mock_ray.remote = Mock() - - with patch('fedgraph.federated_methods.Server') as mock_server_class, \ - patch('fedgraph.federated_methods.torch.manual_seed'): - + + with patch("fedgraph.federated_methods.Server") as mock_server_class, patch( + "fedgraph.federated_methods.torch.manual_seed" + ): mock_server = Mock() mock_server_class.return_value = mock_server - + # Mock remote trainer creation mock_trainer_class = Mock() mock_ray.remote.return_value = mock_trainer_class - + try: run_NC(self.args, self.mock_data) except Exception: - # Expected to fail due to complex initialization, + # Expected to fail due to complex initialization, # but we verify the setup calls pass - + # Verify initialization calls mock_ray.init.assert_called() if self.args.monitor: mock_monitor.assert_called_once() - - @patch('fedgraph.federated_methods.ray') + + @patch("fedgraph.federated_methods.ray") def test_run_nc_without_ray(self, mock_ray): """Test run_NC without Ray distributed computing.""" self.args.use_ray = False mock_ray.init = Mock() - - with patch('fedgraph.federated_methods.Server') as mock_server_class, \ - patch('fedgraph.federated_methods.Trainer_General') as mock_trainer_class, \ - patch('fedgraph.federated_methods.torch.manual_seed'): - + + with patch("fedgraph.federated_methods.Server") as mock_server_class, patch( + "fedgraph.federated_methods.Trainer_General" + ) as mock_trainer_class, patch("fedgraph.federated_methods.torch.manual_seed"): mock_server = Mock() mock_server_class.return_value = mock_server mock_server.trainers = [] - + mock_trainer = Mock() mock_trainer_class.return_value = mock_trainer - + try: run_NC(self.args, self.mock_data) except Exception: # Expected to fail due to complex flow, but verify no Ray init pass - + mock_ray.init.assert_not_called() class TestRunGC: """Test run_GC function for graph classification.""" - + def setup_method(self): """Setup test data for GC.""" self.args = attridict.AttriDict() @@ -326,72 +341,76 @@ def setup_method(self): self.args.device = "cpu" self.args.monitor = True self.args.seed = 42 - + # Mock data for GC self.mock_data = { "0-PROTEINS": ( {"train": Mock(), "val": Mock(), "test": Mock()}, 10, # num_node_features - 2, # num_graph_labels - 100 # train_size + 2, # num_graph_labels + 100, # train_size ), "1-PROTEINS": ( {"train": Mock(), "val": Mock(), "test": Mock()}, 10, # num_node_features - 2, # num_graph_labels - 100 # train_size - ) + 2, # num_graph_labels + 100, # train_size + ), } - - @patch('fedgraph.federated_methods.setup_server') - @patch('fedgraph.federated_methods.setup_trainers') - @patch('fedgraph.federated_methods.Monitor') - def test_run_gc_basic_setup(self, mock_monitor, mock_setup_trainers, mock_setup_server): + + @patch("fedgraph.federated_methods.setup_server") + @patch("fedgraph.federated_methods.setup_trainers") + @patch("fedgraph.federated_methods.Monitor") + def test_run_gc_basic_setup( + self, mock_monitor, mock_setup_trainers, mock_setup_server + ): """Test basic setup of run_GC function.""" mock_monitor_instance = Mock() mock_monitor.return_value = mock_monitor_instance - + mock_server = Mock() mock_setup_server.return_value = mock_server - + mock_trainers = [Mock(), Mock()] mock_setup_trainers.return_value = mock_trainers - - with patch('fedgraph.federated_methods.run_GC_Fed_algorithm') as mock_run_fed: + + with patch("fedgraph.federated_methods.run_GC_Fed_algorithm") as mock_run_fed: run_GC(self.args, self.mock_data) - + mock_setup_server.assert_called_once() mock_setup_trainers.assert_called_once() mock_run_fed.assert_called_once() - - @patch('fedgraph.federated_methods.run_GC_selftrain') + + @patch("fedgraph.federated_methods.run_GC_selftrain") def test_run_gc_selftrain(self, mock_run_selftrain): """Test run_GC with SelfTrain method.""" self.args.method = "SelfTrain" - + run_GC(self.args, self.mock_data) - + mock_run_selftrain.assert_called_once_with(self.args, self.mock_data) - - @patch('fedgraph.federated_methods.run_GCFL_algorithm') + + @patch("fedgraph.federated_methods.run_GCFL_algorithm") def test_run_gc_gcfl(self, mock_run_gcfl): """Test run_GC with GCFL method.""" self.args.method = "GCFL" - - with patch('fedgraph.federated_methods.setup_server') as mock_setup_server, \ - patch('fedgraph.federated_methods.setup_trainers') as mock_setup_trainers: - + + with patch( + "fedgraph.federated_methods.setup_server" + ) as mock_setup_server, patch( + "fedgraph.federated_methods.setup_trainers" + ) as mock_setup_trainers: mock_setup_server.return_value = Mock() mock_setup_trainers.return_value = [Mock(), Mock()] - + run_GC(self.args, self.mock_data) - + mock_run_gcfl.assert_called_once() class TestRunLP: """Test run_LP function for link prediction.""" - + def setup_method(self): """Setup test data for LP.""" self.args = Mock() @@ -401,56 +420,59 @@ def setup_method(self): self.args.device = "cpu" self.args.monitor = True self.args.seed = 42 - - @patch('fedgraph.federated_methods.check_data_files_existance') - @patch('fedgraph.federated_methods.get_global_user_item_mapping') - @patch('fedgraph.federated_methods.get_start_end_time') - def test_run_lp_basic_setup(self, mock_get_time, mock_get_mapping, mock_check_files): + + @patch("fedgraph.federated_methods.check_data_files_existance") + @patch("fedgraph.federated_methods.get_global_user_item_mapping") + @patch("fedgraph.federated_methods.get_start_end_time") + def test_run_lp_basic_setup( + self, mock_get_time, mock_get_mapping, mock_check_files + ): """Test basic setup of run_LP function.""" mock_check_files.return_value = True mock_get_mapping.return_value = (Mock(), Mock()) mock_get_time.return_value = (0, 100) - - with patch('fedgraph.federated_methods.Server_LP') as mock_server_class, \ - patch('fedgraph.federated_methods.Monitor') as mock_monitor, \ - patch('fedgraph.federated_methods.ray'): - + + with patch("fedgraph.federated_methods.Server_LP") as mock_server_class, patch( + "fedgraph.federated_methods.Monitor" + ) as mock_monitor, patch("fedgraph.federated_methods.ray"): mock_server = Mock() mock_server_class.return_value = mock_server mock_monitor_instance = Mock() mock_monitor.return_value = mock_monitor_instance - + try: run_LP(self.args) except Exception: # Expected to fail due to complex initialization pass - + # Verify setup calls mock_check_files.assert_called_once() mock_get_mapping.assert_called_once() mock_get_time.assert_called_once() - - @patch('fedgraph.federated_methods.check_data_files_existance') + + @patch("fedgraph.federated_methods.check_data_files_existance") def test_run_lp_missing_data_files(self, mock_check_files): """Test run_LP when data files are missing.""" mock_check_files.return_value = False - + with pytest.raises(FileNotFoundError): run_LP(self.args) class TestFederatedMethodsIntegration: """Integration tests for federated methods.""" - - @patch('fedgraph.federated_methods.data_loader') - @patch('fedgraph.federated_methods.run_NC') - @patch('fedgraph.federated_methods.run_GC') - @patch('fedgraph.federated_methods.run_LP') - def test_task_routing(self, mock_run_lp, mock_run_gc, mock_run_nc, mock_data_loader): + + @patch("fedgraph.federated_methods.data_loader") + @patch("fedgraph.federated_methods.run_NC") + @patch("fedgraph.federated_methods.run_GC") + @patch("fedgraph.federated_methods.run_LP") + def test_task_routing( + self, mock_run_lp, mock_run_gc, mock_run_nc, mock_data_loader + ): """Test that run_fedgraph routes to correct task functions.""" mock_data_loader.return_value = Mock() - + # Test NC routing args_nc = attridict.AttriDict() args_nc.fedgraph_task = "NC" @@ -458,69 +480,69 @@ def test_task_routing(self, mock_run_lp, mock_run_gc, mock_run_nc, mock_data_loa args_nc.use_lowrank = False run_fedgraph(args_nc) mock_run_nc.assert_called_once() - + # Test GC routing args_gc = attridict.AttriDict() args_gc.fedgraph_task = "GC" args_gc.use_huggingface = False run_fedgraph(args_gc) mock_run_gc.assert_called_once() - + # Test LP routing args_lp = attridict.AttriDict() args_lp.fedgraph_task = "LP" run_fedgraph(args_lp) mock_run_lp.assert_called_once() - + def test_argument_validation_comprehensive(self): """Test comprehensive argument validation.""" # Test invalid task args = attridict.AttriDict() args.fedgraph_task = "INVALID" - - with patch('fedgraph.federated_methods.data_loader'): + + with patch("fedgraph.federated_methods.data_loader"): with pytest.raises((KeyError, AttributeError)): run_fedgraph(args) - + # Test low-rank with invalid method args = attridict.AttriDict() args.fedgraph_task = "NC" args.use_lowrank = True args.method = "GCFL" args.use_encryption = False - + with pytest.raises(ValueError): run_fedgraph(args) - - @patch('fedgraph.federated_methods.data_loader') + + @patch("fedgraph.federated_methods.data_loader") def test_data_loading_logic(self, mock_data_loader): """Test data loading logic in different scenarios.""" mock_data = Mock() mock_data_loader.return_value = mock_data - - with patch('fedgraph.federated_methods.run_NC') as mock_run_nc: + + with patch("fedgraph.federated_methods.run_NC") as mock_run_nc: # Test normal data loading args = attridict.AttriDict() args.fedgraph_task = "NC" args.use_huggingface = False args.use_lowrank = False - + run_fedgraph(args) - + mock_data_loader.assert_called_once_with(args) mock_run_nc.assert_called_once_with(args, mock_data) - + # Reset mocks mock_data_loader.reset_mock() - - with patch('fedgraph.federated_methods.run_NC') as mock_run_nc: + + with patch("fedgraph.federated_methods.run_NC") as mock_run_nc: # Test Hugging Face data loading (should skip data_loader) args = attridict.AttriDict() args.fedgraph_task = "NC" args.use_huggingface = True args.use_lowrank = False - + run_fedgraph(args) - + mock_data_loader.assert_not_called() mock_run_nc.assert_called_once_with(args, None) diff --git a/tests/unit/test_server_class.py b/tests/unit/test_server_class.py index 9ad5c18..dd6fdd6 100644 --- a/tests/unit/test_server_class.py +++ b/tests/unit/test_server_class.py @@ -1,21 +1,22 @@ +from unittest.mock import MagicMock, Mock, patch + +import numpy as np import pytest import torch -import numpy as np -from unittest.mock import Mock, patch, MagicMock from fedgraph.server_class import Server, Server_GC class TestServer: """Test Server class for federated learning.""" - + def setup_method(self): """Setup common test data.""" self.feature_dim = 100 self.args_hidden = 64 self.class_num = 7 - self.device = torch.device('cpu') - + self.device = torch.device("cpu") + # Mock trainers self.mock_trainers = [] for i in range(3): @@ -25,139 +26,139 @@ def setup_method(self): torch.randn(64, 100), torch.randn(64), torch.randn(7, 64), - torch.randn(7) + torch.randn(7), ) - trainer.local_node_index = torch.arange(i*10, (i+1)*10) - trainer.communicate_node_index = torch.arange(i*20, (i+1)*20) + trainer.local_node_index = torch.arange(i * 10, (i + 1) * 10) + trainer.communicate_node_index = torch.arange(i * 20, (i + 1) * 20) self.mock_trainers.append(trainer) - + # Mock args self.args = Mock() self.args.num_hops = 2 self.args.dataset = "cora" self.args.num_layers = 2 self.args.method = "FedAvg" - - @patch('fedgraph.server_class.AggreGCN') + + @patch("fedgraph.server_class.AggreGCN") def test_server_init_with_hops(self, mock_aggre_gcn): """Test Server initialization with FedGCN-style num_hops.""" mock_model = Mock() mock_aggre_gcn.return_value = mock_model mock_model.to.return_value = mock_model - + server = Server( feature_dim=self.feature_dim, args_hidden=self.args_hidden, class_num=self.class_num, device=self.device, trainers=self.mock_trainers, - args=self.args + args=self.args, ) - + assert server.args == self.args assert server.model is not None assert server.trainers == self.mock_trainers assert server.num_of_trainers == len(self.mock_trainers) mock_aggre_gcn.assert_called_once() - - @patch('fedgraph.server_class.GCN_arxiv') + + @patch("fedgraph.server_class.GCN_arxiv") def test_server_init_arxiv_dataset(self, mock_gcn_arxiv): """Test Server initialization with arxiv dataset.""" mock_model = Mock() mock_gcn_arxiv.return_value = mock_model mock_model.to.return_value = mock_model - + self.args.dataset = "ogbn-arxiv" self.args.num_hops = 0 # FedAvg method - + server = Server( feature_dim=self.feature_dim, args_hidden=self.args_hidden, class_num=self.class_num, device=self.device, trainers=self.mock_trainers, - args=self.args + args=self.args, ) - + mock_gcn_arxiv.assert_called_once() - - @patch('fedgraph.server_class.SAGE_products') + + @patch("fedgraph.server_class.SAGE_products") def test_server_init_products_dataset(self, mock_sage): """Test Server initialization with products dataset.""" mock_model = Mock() mock_sage.return_value = mock_model mock_model.to.return_value = mock_model - + self.args.dataset = "ogbn-products" self.args.num_hops = 0 - + server = Server( feature_dim=self.feature_dim, args_hidden=self.args_hidden, class_num=self.class_num, device=self.device, trainers=self.mock_trainers, - args=self.args + args=self.args, ) - + mock_sage.assert_called_once() - + def test_zero_params(self): """Test zero_params method.""" - with patch('fedgraph.server_class.AggreGCN') as mock_gcn: + with patch("fedgraph.server_class.AggreGCN") as mock_gcn: mock_model = Mock() mock_gcn.return_value = mock_model mock_model.to.return_value = mock_model - + # Mock parameters mock_param1 = Mock() mock_param1.data = torch.randn(10, 5) mock_param2 = Mock() mock_param2.data = torch.randn(5) mock_model.parameters.return_value = [mock_param1, mock_param2] - + server = Server( feature_dim=self.feature_dim, args_hidden=self.args_hidden, class_num=self.class_num, device=self.device, trainers=self.mock_trainers, - args=self.args + args=self.args, ) - + server.zero_params() - + # Check that parameters were zeroed mock_model.parameters.assert_called() - + def test_broadcast_params(self): """Test broadcast_params method.""" - with patch('fedgraph.server_class.AggreGCN') as mock_gcn: + with patch("fedgraph.server_class.AggreGCN") as mock_gcn: mock_model = Mock() mock_model.parameters.return_value = [ torch.randn(32, 20, requires_grad=True), - torch.randn(32, requires_grad=True) + torch.randn(32, requires_grad=True), ] mock_gcn.return_value = mock_model mock_model.to.return_value = mock_model - mock_model.state_dict.return_value = {'layer.weight': torch.randn(10, 5)} - + mock_model.state_dict.return_value = {"layer.weight": torch.randn(10, 5)} + server = Server( feature_dim=self.feature_dim, args_hidden=self.args_hidden, class_num=self.class_num, device=self.device, trainers=self.mock_trainers, - args=self.args + args=self.args, ) - + # Mock update_params method for trainers for trainer in server.trainers: trainer.update_params = Mock() - + with patch("fedgraph.server_class.ray.get") as mock_ray_get: server.broadcast_params(current_global_epoch=1) - + # Verify that all trainers received parameter updates for trainer in server.trainers: trainer.update_params.remote.assert_called_once() @@ -193,9 +194,7 @@ def wait_for_one(refs, **_kwargs): with patch("fedgraph.server_class.random.sample", return_value=[0, 2]), patch( "fedgraph.server_class.ray.get", side_effect=resolve - ), patch( - "fedgraph.server_class.ray.wait", side_effect=wait_for_one - ), patch( + ), patch("fedgraph.server_class.ray.wait", side_effect=wait_for_one), patch( "fedgraph.server_class.time.time", side_effect=[10.0, 12.0, 12.0, 15.0] ): round_stats = server.train(4, sample_ratio=0.75) @@ -250,9 +249,7 @@ def resolve(refs): def wait_for_one(refs, **_kwargs): return refs[:1], refs[1:] - with patch( - "fedgraph.server_class.ray.get", side_effect=resolve - ), patch( + with patch("fedgraph.server_class.ray.get", side_effect=resolve), patch( "fedgraph.server_class.ray.wait", side_effect=wait_for_one ), patch( "fedgraph.server_class.time.time", @@ -300,9 +297,11 @@ def serialize(self): "fedgraph.server_class.ts.ckks_vector_from", side_effect=lambda _context, payload: FakeCKKSVector(payload), ), patch("fedgraph.server_class.time.time", side_effect=[10.0, 13.0]): - aggregated_params, metadata, aggregation_time = ( - server.aggregate_encrypted_params(encrypted_params_list) - ) + ( + aggregated_params, + metadata, + aggregation_time, + ) = server.aggregate_encrypted_params(encrypted_params_list) assert aggregated_params == [pytest.approx(3000.0)] assert metadata == [{"shape": torch.Size([1]), "scale": 1000.0}] @@ -365,49 +364,50 @@ def test_aggregate_encrypted_params_rejects_invalid_scale(self): with pytest.raises(ValueError, match="scale must be positive"): server.aggregate_encrypted_params(encrypted_params_list) - + def test_get_model_size(self): """Test get_model_size method.""" - with patch('fedgraph.server_class.AggreGCN') as mock_gcn: + with patch("fedgraph.server_class.AggreGCN") as mock_gcn: mock_model = Mock() mock_gcn.return_value = mock_model mock_model.to.return_value = mock_model - + # Mock parameters with known sizes param1 = torch.randn(10, 5) # 50 elements - param2 = torch.randn(5) # 5 elements + param2 = torch.randn(5) # 5 elements mock_model.parameters.return_value = [param1, param2] - + server = Server( feature_dim=self.feature_dim, args_hidden=self.args_hidden, class_num=self.class_num, device=self.device, trainers=self.mock_trainers, - args=self.args + args=self.args, ) - + model_size = server.get_model_size() - + assert isinstance(model_size, float) assert model_size > 0 - + + class TestServerGC: """Test Server_GC class for graph classification.""" - + def setup_method(self): """Setup common test data for GC server.""" - self.device = torch.device('cpu') + self.device = torch.device("cpu") self.use_cluster = False - + # Mock model self.mock_model = Mock() self.mock_model.to.return_value = self.mock_model self.mock_model.named_parameters.return_value = [ - ('layer1.weight', torch.randn(10, 5, requires_grad=True)), - ('layer1.bias', torch.randn(10, requires_grad=True)) + ("layer1.weight", torch.randn(10, 5, requires_grad=True)), + ("layer1.bias", torch.randn(10, requires_grad=True)), ] - + # Mock trainers for GC self.mock_trainers = [] for i in range(3): @@ -415,130 +415,113 @@ def setup_method(self): trainer.id = i trainer.train_size = 100 trainer.W = { - 'layer1.weight': torch.randn(10, 5, requires_grad=True), - 'layer1.bias': torch.randn(10, requires_grad=True) + "layer1.weight": torch.randn(10, 5, requires_grad=True), + "layer1.bias": torch.randn(10, requires_grad=True), } trainer.compute_update_norm = Mock(return_value=0.5) trainer.compute_mean_norm = Mock(return_value=torch.randn(15)) self.mock_trainers.append(trainer) - + def test_server_gc_init(self): """Test Server_GC initialization.""" server = Server_GC( - model=self.mock_model, - device=self.device, - use_cluster=self.use_cluster + model=self.mock_model, device=self.device, use_cluster=self.use_cluster ) - + assert server.model == self.mock_model assert server.use_cluster == self.use_cluster - assert hasattr(server, 'W') - assert hasattr(server, 'model_cache') - + assert hasattr(server, "W") + assert hasattr(server, "model_cache") + def test_random_sample_trainers(self): """Test random_sample_trainers method.""" server = Server_GC( - model=self.mock_model, - device=self.device, - use_cluster=self.use_cluster + model=self.mock_model, device=self.device, use_cluster=self.use_cluster ) - + all_trainers = self.mock_trainers frac = 0.6 # Select 60% of trainers - + selected = server.random_sample_trainers(all_trainers, frac) - + assert isinstance(selected, list) assert len(selected) <= len(all_trainers) assert len(selected) >= 1 # At least one trainer should be selected for trainer in selected: assert trainer in all_trainers - + def test_aggregate_weights(self): """Test aggregate_weights method.""" server = Server_GC( - model=self.mock_model, - device=self.device, - use_cluster=self.use_cluster + model=self.mock_model, device=self.device, use_cluster=self.use_cluster ) - + selected_trainers = self.mock_trainers[:2] # Select first 2 trainers - + server.aggregate_weights(selected_trainers) - + # Verify that the method completes without error # The actual aggregation logic depends on the specific implementation assert True # Method executed successfully - + def test_compute_pairwise_similarities(self): """Test compute_pairwise_similarities method.""" server = Server_GC( - model=self.mock_model, - device=self.device, - use_cluster=self.use_cluster + model=self.mock_model, device=self.device, use_cluster=self.use_cluster ) - + # Ensure trainers have consistent weight shapes for trainer in self.mock_trainers: trainer.W = { - 'layer1.weight': torch.randn(10, 5), - 'layer1.bias': torch.randn(10) + "layer1.weight": torch.randn(10, 5), + "layer1.bias": torch.randn(10), } - + similarities = server.compute_pairwise_similarities(self.mock_trainers) - + assert isinstance(similarities, np.ndarray) assert similarities.shape == (len(self.mock_trainers), len(self.mock_trainers)) # Check symmetry assert np.allclose(similarities, similarities.T) # Check diagonal elements are 1 (self-similarity) assert np.allclose(np.diag(similarities), 1.0) - + def test_compute_pairwise_distances(self): """Test compute_pairwise_distances method.""" server = Server_GC( - model=self.mock_model, - device=self.device, - use_cluster=self.use_cluster + model=self.mock_model, device=self.device, use_cluster=self.use_cluster ) - + # Create gradient norm sequences for DTW computation gradient_sequences = [ [0.9, 0.8, 0.7, 0.6], # Sequence for trainer 0 - [0.8, 0.7, 0.6, 0.5], # Sequence for trainer 1 - [0.7, 0.6, 0.5, 0.4] # Sequence for trainer 2 + [0.8, 0.7, 0.6, 0.5], # Sequence for trainer 1 + [0.7, 0.6, 0.5, 0.4], # Sequence for trainer 2 ] - + distances = server.compute_pairwise_distances( - gradient_sequences, - standardize=False + gradient_sequences, standardize=False ) - + assert isinstance(distances, np.ndarray) assert distances.shape == (len(gradient_sequences), len(gradient_sequences)) # Distance matrix should be symmetric assert np.allclose(distances, distances.T) # Diagonal should be zero (distance to self) assert np.allclose(np.diag(distances), 0.0) - + def test_min_cut(self): """Test min_cut method.""" server = Server_GC( - model=self.mock_model, - device=self.device, - use_cluster=self.use_cluster + model=self.mock_model, device=self.device, use_cluster=self.use_cluster ) - + # Create a simple similarity matrix - similarity = np.array([ - [1.0, 0.8, 0.3], - [0.8, 1.0, 0.4], - [0.3, 0.4, 1.0] - ]) + similarity = np.array([[1.0, 0.8, 0.3], [0.8, 1.0, 0.4], [0.3, 0.4, 1.0]]) idc = [0, 1, 2] - + cluster1, cluster2 = server.min_cut(similarity, idc) - + assert isinstance(cluster1, list) assert isinstance(cluster2, list) assert len(cluster1) + len(cluster2) == len(idc) @@ -546,41 +529,37 @@ def test_min_cut(self): assert set(cluster1).isdisjoint(set(cluster2)) # Ensure all indices are covered assert set(cluster1).union(set(cluster2)) == set(idc) - + def test_aggregate_clusterwise(self): """Test aggregate_clusterwise method.""" server = Server_GC( - model=self.mock_model, - device=self.device, - use_cluster=self.use_cluster + model=self.mock_model, device=self.device, use_cluster=self.use_cluster ) - + # Create trainer clusters trainer_clusters = [ [self.mock_trainers[0], self.mock_trainers[1]], - [self.mock_trainers[2]] + [self.mock_trainers[2]], ] - + server.aggregate_clusterwise(trainer_clusters) - + # Verify method execution assert True # Method completed without error - + def test_compute_max_update_norm(self): """Test compute_max_update_norm method.""" server = Server_GC( - model=self.mock_model, - device=self.device, - use_cluster=self.use_cluster + model=self.mock_model, device=self.device, use_cluster=self.use_cluster ) - + cluster = self.mock_trainers[:2] - + max_norm = server.compute_max_update_norm(cluster) - + assert isinstance(max_norm, float) assert max_norm >= 0 - + # Verify that compute_update_norm was called on trainers for trainer in cluster: trainer.compute_update_norm.assert_called() @@ -588,47 +567,44 @@ def test_compute_max_update_norm(self): class TestServerIntegration: """Integration tests for server classes.""" - - @patch('fedgraph.server_class.AggreGCN') + + @patch("fedgraph.server_class.AggreGCN") def test_server_trainer_interaction(self, mock_gcn_class): """Test interaction between server and trainers.""" # Setup feature_dim = 50 args_hidden = 32 class_num = 3 - device = torch.device('cpu') - + device = torch.device("cpu") + args = Mock() args.num_hops = 2 args.dataset = "cora" args.num_layers = 2 args.method = "FedAvg" - + # Mock model mock_model = Mock() mock_gcn_class.return_value = mock_model mock_model.to.return_value = mock_model mock_model.state_dict.return_value = { - 'layer1.weight': torch.randn(32, 50), - 'layer1.bias': torch.randn(32) + "layer1.weight": torch.randn(32, 50), + "layer1.bias": torch.randn(32), } mock_model.parameters.return_value = [ torch.randn(32, 50), torch.randn(32), ] - + # Create mock trainers trainers = [] for i in range(2): trainer = Mock() trainer.rank = i trainer.update_params = Mock() - trainer.get_params.return_value = ( - torch.randn(32, 50), - torch.randn(32) - ) + trainer.get_params.return_value = (torch.randn(32, 50), torch.randn(32)) trainers.append(trainer) - + # Create server server = Server( feature_dim=feature_dim, @@ -636,36 +612,36 @@ def test_server_trainer_interaction(self, mock_gcn_class): class_num=class_num, device=device, trainers=trainers, - args=args + args=args, ) - + # Test parameter broadcast with patch("fedgraph.server_class.ray.get"): server.broadcast_params(current_global_epoch=1) - + # Verify all trainers received updates for trainer in trainers: trainer.update_params.remote.assert_called_once() - + # Test model size computation model_size = server.get_model_size() assert isinstance(model_size, float) assert model_size > 0 - + def test_server_gc_clustering_workflow(self): """Test complete clustering workflow in Server_GC.""" - device = torch.device('cpu') + device = torch.device("cpu") use_cluster = True - + # Mock model mock_model = Mock() mock_model.to.return_value = mock_model mock_model.named_parameters.return_value = [ - ('layer1.weight', torch.randn(10, 5, requires_grad=True)), + ("layer1.weight", torch.randn(10, 5, requires_grad=True)), ] - + server = Server_GC(model=mock_model, device=device, use_cluster=use_cluster) - + # Create trainers with consistent data trainers = [] for i in range(4): @@ -673,35 +649,40 @@ def test_server_gc_clustering_workflow(self): trainer.id = i trainer.train_size = 100 trainer.W = { - 'layer1.weight': torch.randn(10, 5), - 'layer1.bias': torch.randn(10) + "layer1.weight": torch.randn(10, 5), + "layer1.bias": torch.randn(10), } trainer.train_stats = { - 'trainingLosses': [0.9 - i*0.1, 0.8 - i*0.1], - 'trainingAccs': [0.6 + i*0.1, 0.7 + i*0.1] + "trainingLosses": [0.9 - i * 0.1, 0.8 - i * 0.1], + "trainingAccs": [0.6 + i * 0.1, 0.7 + i * 0.1], } - trainer.compute_update_norm = Mock(return_value=0.5 + i*0.1) + trainer.compute_update_norm = Mock(return_value=0.5 + i * 0.1) trainer.compute_mean_norm = Mock(return_value=torch.randn(15)) trainers.append(trainer) - + # Test similarity computation similarities = server.compute_pairwise_similarities(trainers) assert similarities.shape == (4, 4) - + # Test distance computation with gradient sequences gradient_sequences = [ - [0.9 - i*0.1, 0.8 - i*0.1, 0.7 - i*0.1] for i in range(4) + [0.9 - i * 0.1, 0.8 - i * 0.1, 0.7 - i * 0.1] for i in range(4) ] - distances = server.compute_pairwise_distances(gradient_sequences, standardize=False) + distances = server.compute_pairwise_distances( + gradient_sequences, standardize=False + ) assert distances.shape == (4, 4) - + # Test clustering cluster1, cluster2 = server.min_cut(similarities, list(range(4))) assert len(cluster1) + len(cluster2) == 4 - + # Test clusterwise aggregation - trainer_clusters = [[trainers[i] for i in cluster1], [trainers[i] for i in cluster2]] + trainer_clusters = [ + [trainers[i] for i in cluster1], + [trainers[i] for i in cluster2], + ] server.aggregate_clusterwise(trainer_clusters) - + # Verify workflow completed successfully assert True diff --git a/tests/unit/test_trainer_class.py b/tests/unit/test_trainer_class.py index af94659..00cc451 100644 --- a/tests/unit/test_trainer_class.py +++ b/tests/unit/test_trainer_class.py @@ -1,69 +1,70 @@ -import pytest -import torch +from unittest.mock import MagicMock, Mock, patch + import numpy as np -from unittest.mock import Mock, patch, MagicMock +import pytest import ray +import torch from huggingface_hub.errors import EntryNotFoundError from fedgraph.trainer_class import ( - Trainer_General, Trainer_GC, - load_trainer_data_from_hugging_face + Trainer_General, + load_trainer_data_from_hugging_face, ) class TestLoadTrainerDataFromHuggingFace: """Test load_trainer_data_from_hugging_face function.""" - - @patch('fedgraph.trainer_class.hf_hub_download') - @patch('builtins.open') - @patch('torch.load') - def test_load_trainer_data_success(self, mock_torch_load, mock_open, mock_hf_download): + + @patch("fedgraph.trainer_class.hf_hub_download") + @patch("builtins.open") + @patch("torch.load") + def test_load_trainer_data_success( + self, mock_torch_load, mock_open, mock_hf_download + ): """Test successful loading of trainer data from Hugging Face.""" # Setup mocks mock_hf_download.return_value = "/tmp/test_file.pt" mock_file = Mock() mock_file.read.return_value = b"test_tensor_data" mock_open.return_value.__enter__.return_value = mock_file - + # Mock tensor data mock_tensors = [ torch.randn(100), # local_node_index - torch.randn(50), # communicate_node_global_index + torch.randn(50), # communicate_node_global_index torch.randn(2, 200), # global_edge_index_client - torch.randn(80), # train_labels - torch.randn(20), # test_labels + torch.randn(80), # train_labels + torch.randn(20), # test_labels torch.randn(100, 10), # features - torch.randn(80), # in_com_train_node_local_indexes - torch.randn(20), # in_com_test_node_local_indexes + torch.randn(80), # in_com_train_node_local_indexes + torch.randn(20), # in_com_test_node_local_indexes torch.tensor(100), # global_node_num torch.tensor(3), # class_num ] mock_torch_load.side_effect = mock_tensors - + args = Mock() args.dataset = "cora" args.n_trainer = 5 args.num_hops = 2 args.iid_beta = 0.5 - + result = load_trainer_data_from_hugging_face(trainer_id=0, args=args) - + assert len(result) == 10 assert all(isinstance(tensor, torch.Tensor) for tensor in result) - + # Verify calls assert mock_hf_download.call_count == 10 expected_repo = "FedGraph/fedgraph_cora_5trainer_2hop_iid_beta_0.5_trainer_id_0" mock_hf_download.assert_any_call( - repo_id=expected_repo, - repo_type="dataset", - filename="local_node_index.pt" + repo_id=expected_repo, repo_type="dataset", filename="local_node_index.pt" ) - @patch('fedgraph.trainer_class.hf_hub_download') - @patch('builtins.open') - @patch('torch.load') + @patch("fedgraph.trainer_class.hf_hub_download") + @patch("builtins.open") + @patch("torch.load") def test_load_existing_repo_without_global_metadata( self, mock_torch_load, mock_open, mock_hf_download ): @@ -89,13 +90,13 @@ def download_side_effect(*, filename, **kwargs): class TestTrainerGeneral: """Test Trainer_General class.""" - + def setup_method(self): """Setup common test data.""" self.rank = 0 - self.device = torch.device('cpu') + self.device = torch.device("cpu") self.args_hidden = 64 - + # Create mock data self.local_node_index = torch.arange(0, 50) self.communicate_node_index = torch.arange(0, 100) @@ -105,7 +106,7 @@ def setup_method(self): self.features = torch.randn(100, 50) self.idx_train = torch.arange(0, 40) self.idx_test = torch.arange(40, 50) - + self.args = Mock() self.args.local_step = 5 self.args.method = "FedAvg" @@ -114,8 +115,8 @@ def setup_method(self): self.args.learning_rate = 0.01 self.args.dataset = "cora" self.args.batch_size = 0 # No batching for simplicity - - @patch('fedgraph.trainer_class.load_trainer_data_from_hugging_face') + + @patch("fedgraph.trainer_class.load_trainer_data_from_hugging_face") def test_trainer_init_with_data(self, mock_load_data): """Test Trainer_General initialization with provided data.""" trainer = Trainer_General( @@ -134,18 +135,20 @@ def test_trainer_init_with_data(self, mock_load_data): global_node_num=100, class_num=3, ) - + assert trainer.rank == self.rank assert trainer.device == self.device assert trainer.args_hidden == self.args_hidden assert trainer.local_step == self.args.local_step assert trainer.global_node_num == 100 assert trainer.class_num == 3 - assert torch.equal(trainer.local_node_index, self.local_node_index.to(self.device)) + assert torch.equal( + trainer.local_node_index, self.local_node_index.to(self.device) + ) assert trainer.feature_aggregation is not None # Should be set for FedAvg mock_load_data.assert_not_called() - - @patch('fedgraph.trainer_class.load_trainer_data_from_hugging_face') + + @patch("fedgraph.trainer_class.load_trainer_data_from_hugging_face") def test_trainer_init_without_data(self, mock_load_data): """Test Trainer_General initialization without provided data.""" mock_load_data.return_value = ( @@ -160,19 +163,19 @@ def test_trainer_init_without_data(self, mock_load_data): torch.tensor(100), torch.tensor(3), ) - + trainer = Trainer_General( rank=self.rank, args_hidden=self.args_hidden, device=self.device, - args=self.args + args=self.args, ) - + assert trainer.rank == self.rank assert trainer.global_node_num == 100 assert trainer.class_num == 3 mock_load_data.assert_called_once_with(self.rank, self.args) - + def test_get_info(self): """Test get_info method.""" trainer = Trainer_General( @@ -187,15 +190,17 @@ def test_get_info(self): test_labels=self.test_labels, features=self.features, idx_train=self.idx_train, - idx_test=self.idx_test + idx_test=self.idx_test, ) - + info = trainer.get_info() - + assert "features_num" in info assert "label_num" in info assert info["features_num"] == len(self.features) - expected_label_num = max(self.train_labels.max().item(), self.test_labels.max().item()) + 1 + expected_label_num = ( + max(self.train_labels.max().item(), self.test_labels.max().item()) + 1 + ) assert info["label_num"] == expected_label_num @pytest.mark.parametrize( @@ -229,9 +234,9 @@ def test_get_info_handles_empty_labels( ) assert trainer.get_info()["label_num"] == expected_label_num - - @patch('fedgraph.trainer_class.GCN') - @patch('fedgraph.trainer_class.GCN_arxiv') + + @patch("fedgraph.trainer_class.GCN") + @patch("fedgraph.trainer_class.GCN_arxiv") def test_init_model(self, mock_gcn_arxiv, mock_gcn): """Test init_model method.""" trainer = Trainer_General( @@ -246,19 +251,19 @@ def test_init_model(self, mock_gcn_arxiv, mock_gcn): test_labels=self.test_labels, features=self.features, idx_train=self.idx_train, - idx_test=self.idx_test + idx_test=self.idx_test, ) - + mock_model = Mock() mock_gcn.return_value = mock_model self.args.dataset = "cora" self.args.num_hops = 2 - + trainer.init_model(global_node_num=2708, class_num=7) - + assert trainer.model is not None mock_gcn.assert_called_once() - + def test_update_params(self): """Test update_params method.""" trainer = Trainer_General( @@ -273,23 +278,23 @@ def test_update_params(self): test_labels=self.test_labels, features=self.features, idx_train=self.idx_train, - idx_test=self.idx_test + idx_test=self.idx_test, ) - + # Create mock model mock_model = Mock() mock_optimizer = Mock() trainer.model = mock_model trainer.optimizer = mock_optimizer - + # Create mock parameters mock_params = (torch.randn(10, 5), torch.randn(5)) - + trainer.update_params(mock_params, current_global_epoch=1) - + # Verify that model state_dict loading was attempted - assert mock_model.load_state_dict.called or hasattr(trainer, 'model') - + assert mock_model.load_state_dict.called or hasattr(trainer, "model") + def test_get_local_feature_sum(self): """Test get_local_feature_sum method.""" local_features = self.features[self.local_node_index] @@ -305,16 +310,16 @@ def test_get_local_feature_sum(self): test_labels=self.test_labels, features=local_features, idx_train=self.idx_train, - idx_test=self.idx_test + idx_test=self.idx_test, ) trainer.global_node_num = len(self.features) # Mock the get_1hop_feature_sum function - with patch('fedgraph.trainer_class.get_1hop_feature_sum') as mock_get_1hop: + with patch("fedgraph.trainer_class.get_1hop_feature_sum") as mock_get_1hop: mock_get_1hop.return_value = torch.randn(100, 50) - + result = trainer.get_local_feature_sum() - + assert isinstance(result, torch.Tensor) mock_get_1hop.assert_called_once() feature_matrix = mock_get_1hop.call_args.args[0] @@ -341,7 +346,7 @@ def test_get_local_feature_sum_requires_global_node_num(self): with pytest.raises(RuntimeError, match="metadata must be initialized"): trainer.get_local_feature_sum() - + def test_load_feature_aggregation(self): """Test load_feature_aggregation method.""" trainer = Trainer_General( @@ -356,14 +361,14 @@ def test_load_feature_aggregation(self): test_labels=self.test_labels, features=self.features, idx_train=self.idx_train, - idx_test=self.idx_test + idx_test=self.idx_test, ) - + new_features = torch.randn(100, 50) trainer.load_feature_aggregation(new_features) - + assert torch.equal(trainer.feature_aggregation, new_features.to(trainer.device)) - + def test_get_params(self): """Test get_params method.""" trainer = Trainer_General( @@ -378,22 +383,25 @@ def test_get_params(self): test_labels=self.test_labels, features=self.features, idx_train=self.idx_train, - idx_test=self.idx_test + idx_test=self.idx_test, ) - + # Create mock model with state_dict mock_model = Mock() - mock_state_dict = {'layer1.weight': torch.randn(10, 5), 'layer1.bias': torch.randn(10)} + mock_state_dict = { + "layer1.weight": torch.randn(10, 5), + "layer1.bias": torch.randn(10), + } mock_model.state_dict.return_value = mock_state_dict trainer.model = mock_model - + params = trainer.get_params() - + assert isinstance(params, tuple) mock_model.state_dict.assert_called_once() - - @patch('fedgraph.trainer_class.test') - @patch('fedgraph.trainer_class.train') + + @patch("fedgraph.trainer_class.test") + @patch("fedgraph.trainer_class.train") def test_train_method(self, mock_train_func, mock_test_func): """Test train method.""" trainer = Trainer_General( @@ -408,31 +416,31 @@ def test_train_method(self, mock_train_func, mock_test_func): test_labels=self.test_labels, features=self.features, idx_train=self.idx_train, - idx_test=self.idx_test + idx_test=self.idx_test, ) - + # Setup mock model and optimizer mock_model = Mock() mock_optimizer = Mock() trainer.model = mock_model trainer.optimizer = mock_optimizer trainer.class_num = 7 # Add missing class_num - + self.args.batch_size = 0 # Ensure no batching for this test - + mock_train_func.return_value = (0.5, 0.85) # loss, accuracy mock_test_func.return_value = (0.3, 0.9) # loss, accuracy - + trainer.train(current_global_round=1) - + assert mock_train_func.call_count == trainer.local_step assert mock_test_func.call_count == trainer.local_step assert len(trainer.train_losses) == trainer.local_step assert len(trainer.train_accs) == trainer.local_step assert len(trainer.test_losses) == trainer.local_step assert len(trainer.test_accs) == trainer.local_step - - @patch('fedgraph.trainer_class.test') + + @patch("fedgraph.trainer_class.test") def test_local_test(self, mock_test_func): """Test local_test method.""" trainer = Trainer_General( @@ -447,23 +455,23 @@ def test_local_test(self, mock_test_func): test_labels=self.test_labels, features=self.features, idx_train=self.idx_train, - idx_test=self.idx_test + idx_test=self.idx_test, ) - + # Setup mock model mock_model = Mock() trainer.model = mock_model - + mock_test_func.return_value = (0.3, 0.9) # loss, accuracy - + result = trainer.local_test() - + mock_test_func.assert_called() assert isinstance(result, list) assert len(result) == 2 assert len(trainer.test_losses) > 0 assert len(trainer.test_accs) > 0 - + def test_get_rank(self): """Test get_rank method.""" trainer = Trainer_General( @@ -478,11 +486,11 @@ def test_get_rank(self): test_labels=self.test_labels, features=self.features, idx_train=self.idx_train, - idx_test=self.idx_test + idx_test=self.idx_test, ) - + assert trainer.get_rank() == self.rank - + def test_get_all_loss_accuracy(self): """Test get_all_loss_accuray method.""" trainer = Trainer_General( @@ -497,52 +505,48 @@ def test_get_all_loss_accuracy(self): test_labels=self.test_labels, features=self.features, idx_train=self.idx_train, - idx_test=self.idx_test + idx_test=self.idx_test, ) - + # Add some mock data trainer.train_losses = [0.5, 0.4, 0.3] trainer.train_accs = [0.8, 0.85, 0.9] trainer.test_losses = [0.6, 0.5, 0.4] trainer.test_accs = [0.75, 0.8, 0.85] - + result = trainer.get_all_loss_accuray() - + assert isinstance(result, list) assert len(result) == 4 # train_losses, train_accs, test_losses, test_accs class TestTrainerGC: """Test Trainer_GC class.""" - + def setup_method(self): """Setup common test data for GC.""" self.trainer_id = 0 self.trainer_name = "trainer_0" self.train_size = 100 self.args = Mock() - self.args.device = torch.device('cpu') + self.args.device = torch.device("cpu") self.args.local_step = 5 self.args.n_trainer = 3 - + # Mock model with named_parameters self.mock_model = Mock() self.mock_model.named_parameters.return_value = [ - ('layer1.weight', torch.randn(10, 5, requires_grad=True)), - ('layer1.bias', torch.randn(10, requires_grad=True)) + ("layer1.weight", torch.randn(10, 5, requires_grad=True)), + ("layer1.bias", torch.randn(10, requires_grad=True)), ] self.mock_model.to.return_value = self.mock_model - + # Mock dataloader - self.mock_dataloader = { - 'train': Mock(), - 'val': Mock(), - 'test': Mock() - } - + self.mock_dataloader = {"train": Mock(), "val": Mock(), "test": Mock()} + # Mock optimizer self.mock_optimizer = Mock() - + def test_trainer_gc_init(self): """Test Trainer_GC initialization.""" trainer = Trainer_GC( @@ -552,19 +556,19 @@ def test_trainer_gc_init(self): train_size=self.train_size, dataloader=self.mock_dataloader, optimizer=self.mock_optimizer, - args=self.args + args=self.args, ) - + assert trainer.id == self.trainer_id assert trainer.name == self.trainer_name assert trainer.train_size == self.train_size assert trainer.dataloader == self.mock_dataloader assert trainer.optimizer == self.mock_optimizer assert trainer.args == self.args - assert hasattr(trainer, 'W') - assert hasattr(trainer, 'dW') - assert hasattr(trainer, 'W_old') - + assert hasattr(trainer, "W") + assert hasattr(trainer, "dW") + assert hasattr(trainer, "W_old") + def test_update_params_gc(self): """Test update_params method for GC trainer.""" trainer = Trainer_GC( @@ -574,24 +578,24 @@ def test_update_params_gc(self): train_size=self.train_size, dataloader=self.mock_dataloader, optimizer=self.mock_optimizer, - args=self.args + args=self.args, ) - + mock_params = { - 'layer1.weight': torch.randn(10, 5), - 'layer1.bias': torch.randn(10) + "layer1.weight": torch.randn(10, 5), + "layer1.bias": torch.randn(10), } - + trainer.update_params(mock_params) - + # Check that gconv_names is set assert trainer.gconv_names == mock_params.keys() - + # Check that weights are updated for key in mock_params: if key in trainer.W: assert torch.equal(trainer.W[key].data, mock_params[key].data) - + def test_reset_params_gc(self): """Test reset_params method for GC trainer.""" trainer = Trainer_GC( @@ -601,27 +605,27 @@ def test_reset_params_gc(self): train_size=self.train_size, dataloader=self.mock_dataloader, optimizer=self.mock_optimizer, - args=self.args + args=self.args, ) - + # Set gconv_names first - trainer.gconv_names = ['layer1.weight', 'layer1.bias'] - + trainer.gconv_names = ["layer1.weight", "layer1.bias"] + # Store original weights original_weights = {k: v.data.clone() for k, v in trainer.W.items()} - + # Modify current weights for key in trainer.W: trainer.W[key].data = torch.randn_like(trainer.W[key].data) - + # Reset params trainer.reset_params() - + # Check that weights are reset to W_old for gconv layers for key in trainer.gconv_names: if key in trainer.W: assert torch.equal(trainer.W[key].data, trainer.W_old[key].data) - + def test_cache_weights_gc(self): """Test cache_weights method for GC trainer.""" trainer = Trainer_GC( @@ -631,20 +635,20 @@ def test_cache_weights_gc(self): train_size=self.train_size, dataloader=self.mock_dataloader, optimizer=self.mock_optimizer, - args=self.args + args=self.args, ) - + # Modify current weights for key in trainer.W: trainer.W[key].data = torch.randn_like(trainer.W[key].data) - + # Cache weights trainer.cache_weights() - + # Check that W_old is updated with current W values for key in trainer.W: assert torch.equal(trainer.W_old[key].data, trainer.W[key].data) - + def test_compute_update_norm(self): """Test compute_update_norm method.""" trainer = Trainer_GC( @@ -654,20 +658,20 @@ def test_compute_update_norm(self): train_size=self.train_size, dataloader=self.mock_dataloader, optimizer=self.mock_optimizer, - args=self.args + args=self.args, ) - + # Set some gradients for key in trainer.dW: trainer.dW[key] = torch.randn_like(trainer.dW[key]) - + # Pass the actual dict or a subset of keys keys = {k: None for k in trainer.dW.keys()} # Create dict with keys norm = trainer.compute_update_norm(keys) - + assert isinstance(norm, float) assert norm >= 0 - + def test_compute_mean_norm(self): """Test compute_mean_norm method.""" trainer = Trainer_GC( @@ -677,32 +681,32 @@ def test_compute_mean_norm(self): train_size=self.train_size, dataloader=self.mock_dataloader, optimizer=self.mock_optimizer, - args=self.args + args=self.args, ) - + # Set some gradients for key in trainer.dW: trainer.dW[key] = torch.randn_like(trainer.dW[key]) - + # Pass the actual dict or a subset of keys keys = {k: None for k in trainer.dW.keys()} # Create dict with keys total_size = 500 mean_norm = trainer.compute_mean_norm(total_size, keys) - + assert isinstance(mean_norm, torch.Tensor) class TestTrainerIntegration: """Integration tests for trainer classes.""" - - @patch('fedgraph.trainer_class.GCN') + + @patch("fedgraph.trainer_class.GCN") def test_trainer_general_full_workflow(self, mock_gcn_class): """Test full workflow of Trainer_General.""" # Setup rank = 0 - device = torch.device('cpu') + device = torch.device("cpu") args_hidden = 64 - + local_node_index = torch.arange(0, 50) communicate_node_index = torch.arange(0, 100) adj = torch.randn(2, 200) @@ -711,7 +715,7 @@ def test_trainer_general_full_workflow(self, mock_gcn_class): features = torch.randn(100, 50) idx_train = torch.arange(0, 40) idx_test = torch.arange(40, 50) - + args = Mock() args.local_step = 5 args.method = "FedAvg" @@ -720,7 +724,7 @@ def test_trainer_general_full_workflow(self, mock_gcn_class): args.num_layers = 2 args.learning_rate = 0.01 args.batch_size = 0 - + # Create trainer trainer = Trainer_General( rank=rank, @@ -734,31 +738,31 @@ def test_trainer_general_full_workflow(self, mock_gcn_class): test_labels=test_labels, features=features, idx_train=idx_train, - idx_test=idx_test + idx_test=idx_test, ) - + # Test initialization assert trainer.rank == rank assert trainer.device == device - + # Test info retrieval info = trainer.get_info() assert "features_num" in info assert "label_num" in info - + # Mock model initialization mock_model = Mock() mock_gcn_class.return_value = mock_model - + trainer.init_model(global_node_num=2708, class_num=7) assert trainer.model is not None - + # Test parameter operations mock_params = (torch.randn(64, 50), torch.randn(64)) trainer.update_params(mock_params, current_global_epoch=1) - + params = trainer.get_params() assert isinstance(params, tuple) - + # Test rank retrieval assert trainer.get_rank() == rank diff --git a/tests/unit/test_utils_gc_lp.py b/tests/unit/test_utils_gc_lp.py index e031342..ac6e9e0 100644 --- a/tests/unit/test_utils_gc_lp.py +++ b/tests/unit/test_utils_gc_lp.py @@ -1,20 +1,22 @@ +from unittest.mock import MagicMock, Mock, patch + +import numpy as np import pytest import torch -import numpy as np -from unittest.mock import Mock, patch, MagicMock from torch_geometric.data import Data from torch_geometric.loader import DataLoader -# Import utils_gc functions +# Import utils_gc functions try: from fedgraph.utils_gc import ( get_max_degree, get_num_graph_labels, get_stats, - split_data, setup_server, - setup_trainers + setup_trainers, + split_data, ) + UTILS_GC_AVAILABLE = True except ImportError: UTILS_GC_AVAILABLE = False @@ -27,8 +29,9 @@ get_data_loaders_per_time_step, get_global_user_item_mapping, get_start_end_time, - to_next_day + to_next_day, ) + UTILS_LP_AVAILABLE = True except ImportError: UTILS_LP_AVAILABLE = False @@ -37,7 +40,7 @@ @pytest.mark.skipif(not UTILS_GC_AVAILABLE, reason="utils_gc not available") class TestUtilsGC: """Test utils_gc functions.""" - + def setup_method(self): """Setup test data for GC utils.""" # Create mock graph data @@ -47,29 +50,29 @@ def setup_method(self): graph.num_nodes = 10 + i graph.num_node_features = 5 graph.y = torch.tensor([i % 2]) # Binary classification - graph.edge_index = torch.randint(0, 10+i, (2, 15+i)) + graph.edge_index = torch.randint(0, 10 + i, (2, 15 + i)) self.mock_graphs.append(graph) - + def test_get_max_degree(self): """Test get_max_degree function.""" # Create mock dataset with graphs of different max degrees mock_dataset = self.mock_graphs - - with patch('torch_geometric.utils.degree') as mock_degree: + + with patch("torch_geometric.utils.degree") as mock_degree: # Mock degree function to return different max degrees mock_degree.side_effect = [ - torch.tensor([1, 2, 3, 4, 5] + [0]*5), # max degree 5 - torch.tensor([2, 3, 4, 3, 2] + [0]*6), # max degree 4 - torch.tensor([1, 1, 1, 1, 1] + [0]*7), # max degree 1 - torch.tensor([3, 5, 2, 4, 6] + [0]*8), # max degree 6 - torch.tensor([2, 2, 2, 2, 2] + [0]*9), # max degree 2 + torch.tensor([1, 2, 3, 4, 5] + [0] * 5), # max degree 5 + torch.tensor([2, 3, 4, 3, 2] + [0] * 6), # max degree 4 + torch.tensor([1, 1, 1, 1, 1] + [0] * 7), # max degree 1 + torch.tensor([3, 5, 2, 4, 6] + [0] * 8), # max degree 6 + torch.tensor([2, 2, 2, 2, 2] + [0] * 9), # max degree 2 ] - + max_degree = get_max_degree(mock_dataset) - + assert isinstance(max_degree, int) assert max_degree == 6 # Should be the maximum across all graphs - + def test_get_num_graph_labels(self): """Test get_num_graph_labels function.""" # Create graphs with different labels @@ -78,96 +81,90 @@ def test_get_num_graph_labels(self): graph = Mock() graph.y = torch.tensor([i % 3]) # Labels 0, 1, 2 graphs.append(graph) - + num_labels = get_num_graph_labels(graphs) - + assert isinstance(num_labels, int) assert num_labels == 3 # Should detect 3 unique labels - + def test_get_stats(self): """Test get_stats function.""" graphs = self.mock_graphs - + stats = get_stats(graphs) - + assert isinstance(stats, dict) - expected_keys = ['num_graphs', 'num_nodes_avg', 'num_edges_avg', 'num_features'] + expected_keys = ["num_graphs", "num_nodes_avg", "num_edges_avg", "num_features"] for key in expected_keys: assert key in stats - - assert stats['num_graphs'] == len(graphs) - assert isinstance(stats['num_nodes_avg'], (int, float)) - assert isinstance(stats['num_edges_avg'], (int, float)) - assert isinstance(stats['num_features'], int) - + + assert stats["num_graphs"] == len(graphs) + assert isinstance(stats["num_nodes_avg"], (int, float)) + assert isinstance(stats["num_edges_avg"], (int, float)) + assert isinstance(stats["num_features"], int) + def test_split_data_with_train_test_sizes(self): """Test split_data with specific train/test sizes.""" data = self.mock_graphs - + train_data, test_data = split_data( data, train_size=0.8, test_size=0.2, shuffle=True, seed=42 ) - + assert len(train_data) + len(test_data) == len(data) assert len(train_data) == int(0.8 * len(data)) assert len(test_data) == len(data) - len(train_data) - + def test_split_data_with_absolute_sizes(self): """Test split_data with absolute sizes.""" data = self.mock_graphs - + train_data, test_data = split_data( data, train_size=3, test_size=2, shuffle=False, seed=42 ) - + assert len(train_data) == 3 assert len(test_data) == 2 - - @patch('fedgraph.utils_gc.Server_GC') - @patch('fedgraph.utils_gc.GIN') + + @patch("fedgraph.utils_gc.Server_GC") + @patch("fedgraph.utils_gc.GIN") def test_setup_server(self, mock_gin, mock_server_gc): """Test setup_server function.""" mock_model = Mock() mock_gin.return_value = mock_model - + mock_server = Mock() mock_server_gc.return_value = mock_server - + args = Mock() args.device = "cpu" - + num_features = 10 num_classes = 2 - + server = setup_server(args, num_features, num_classes) - + mock_gin.assert_called_once() mock_server_gc.assert_called_once() assert server == mock_server - - @patch('fedgraph.utils_gc.Trainer_GC') + + @patch("fedgraph.utils_gc.Trainer_GC") def test_setup_trainers(self, mock_trainer_gc): """Test setup_trainers function.""" mock_trainer = Mock() mock_trainer_gc.return_value = mock_trainer - + args = Mock() args.device = "cpu" - + mock_model = Mock() mock_data = { - "trainer_0": ( - {"train": Mock(), "val": Mock(), "test": Mock()}, - 10, 2, 100 - ), - "trainer_1": ( - {"train": Mock(), "val": Mock(), "test": Mock()}, - 10, 2, 100 - ) + "trainer_0": ({"train": Mock(), "val": Mock(), "test": Mock()}, 10, 2, 100), + "trainer_1": ({"train": Mock(), "val": Mock(), "test": Mock()}, 10, 2, 100), } - + trainers = setup_trainers(mock_model, mock_data, args) - + assert isinstance(trainers, list) assert len(trainers) == len(mock_data) mock_trainer_gc.assert_called() @@ -176,94 +173,91 @@ def test_setup_trainers(self, mock_trainer_gc): @pytest.mark.skipif(not UTILS_LP_AVAILABLE, reason="utils_lp not available") class TestUtilsLP: """Test utils_lp functions.""" - + def test_check_data_files_existance_missing(self): """Test check_data_files_existance with missing files.""" - with patch('os.path.exists') as mock_exists: + with patch("os.path.exists") as mock_exists: mock_exists.return_value = False - + result = check_data_files_existance("fake_dataset", "/fake/path") - + assert result is False - + def test_check_data_files_existance_present(self): """Test check_data_files_existance with present files.""" - with patch('os.path.exists') as mock_exists: + with patch("os.path.exists") as mock_exists: mock_exists.return_value = True - + result = check_data_files_existance("ml-1m", "/fake/path") - + assert result is True - + def test_get_start_end_time(self): """Test get_start_end_time function.""" # Mock data with timestamps mock_data = { "u1.base": Mock(), - "interactions": [(1, 2, 3, 1000), (2, 3, 4, 2000), (3, 4, 5, 500)] + "interactions": [(1, 2, 3, 1000), (2, 3, 4, 2000), (3, 4, 5, 500)], } - - with patch('fedgraph.utils_lp.get_data') as mock_get_data: + + with patch("fedgraph.utils_lp.get_data") as mock_get_data: mock_get_data.return_value = mock_data - + start_time, end_time = get_start_end_time("ml-1m", "/fake/path") - + assert isinstance(start_time, int) assert isinstance(end_time, int) assert start_time <= end_time - + def test_to_next_day(self): """Test to_next_day function.""" # Test with a timestamp (assuming it's in seconds) current_time = 1000000 # Some timestamp - + next_day = to_next_day(current_time) - + assert isinstance(next_day, int) assert next_day > current_time # Should be exactly 24 hours later (86400 seconds) assert next_day - current_time == 86400 - - @patch('fedgraph.utils_lp.get_data') + + @patch("fedgraph.utils_lp.get_data") def test_get_global_user_item_mapping(self, mock_get_data): """Test get_global_user_item_mapping function.""" # Mock interaction data mock_data = { "interactions": [ - (1, 10, 5, 1000), # user 1, item 10, rating 5, time 1000 - (2, 20, 4, 2000), # user 2, item 20, rating 4, time 2000 - (1, 30, 3, 3000), # user 1, item 30, rating 3, time 3000 + (1, 10, 5, 1000), # user 1, item 10, rating 5, time 1000 + (2, 20, 4, 2000), # user 2, item 20, rating 4, time 2000 + (1, 30, 3, 3000), # user 1, item 30, rating 3, time 3000 ] } mock_get_data.return_value = mock_data - + user_mapping, item_mapping = get_global_user_item_mapping("ml-1m", "/fake/path") - + assert isinstance(user_mapping, dict) assert isinstance(item_mapping, dict) - + # Should map original IDs to sequential indices assert len(user_mapping) == 2 # users 1, 2 assert len(item_mapping) == 3 # items 10, 20, 30 - - @patch('fedgraph.utils_lp.get_data') - @patch('fedgraph.utils_lp.create_temporal_split') + + @patch("fedgraph.utils_lp.get_data") + @patch("fedgraph.utils_lp.create_temporal_split") def test_get_data_loaders_per_time_step(self, mock_create_split, mock_get_data): """Test get_data_loaders_per_time_step function.""" # Mock data mock_data = {"interactions": [(1, 2, 3, 1000)]} mock_get_data.return_value = mock_data - + # Mock temporal split - mock_split_data = { - "train": [(1, 2, 3, 1000)], - "test": [(1, 2, 3, 2000)] - } + mock_split_data = {"train": [(1, 2, 3, 1000)], "test": [(1, 2, 3, 2000)]} mock_create_split.return_value = mock_split_data - + user_mapping = {1: 0} item_mapping = {2: 0} - + result = get_data_loaders_per_time_step( dataset="ml-1m", data_path="/fake/path", @@ -272,9 +266,9 @@ def test_get_data_loaders_per_time_step(self, mock_create_split, mock_get_data): item_mapping=item_mapping, num_users=1, num_items=1, - args=Mock() + args=Mock(), ) - + assert isinstance(result, dict) assert "train" in result assert "test" in result @@ -282,7 +276,7 @@ def test_get_data_loaders_per_time_step(self, mock_create_split, mock_get_data): class TestUtilsIntegration: """Integration tests for utility functions.""" - + @pytest.mark.skipif(not UTILS_GC_AVAILABLE, reason="utils_gc not available") def test_gc_workflow(self): """Test complete GC utility workflow.""" @@ -293,48 +287,50 @@ def test_gc_workflow(self): graph.num_nodes = 5 + i graph.num_node_features = 3 graph.y = torch.tensor([i % 2]) - graph.edge_index = torch.randint(0, 5+i, (2, 8+i)) + graph.edge_index = torch.randint(0, 5 + i, (2, 8 + i)) graphs.append(graph) - + # Test stats computation stats = get_stats(graphs) - assert stats['num_graphs'] == 10 - + assert stats["num_graphs"] == 10 + # Test data splitting train_data, test_data = split_data(graphs, train_size=0.8, test_size=0.2) assert len(train_data) + len(test_data) == len(graphs) - + # Test label counting num_labels = get_num_graph_labels(graphs) assert num_labels == 2 # Binary classification - + @pytest.mark.skipif(not UTILS_LP_AVAILABLE, reason="utils_lp not available") def test_lp_data_workflow(self): """Test complete LP utility workflow.""" - with patch('fedgraph.utils_lp.get_data') as mock_get_data: + with patch("fedgraph.utils_lp.get_data") as mock_get_data: # Mock complete interaction data mock_data = { "interactions": [ (1, 10, 5, 1000), - (2, 20, 4, 2000), + (2, 20, 4, 2000), (1, 30, 3, 3000), - (3, 10, 4, 4000) + (3, 10, 4, 4000), ] } mock_get_data.return_value = mock_data - + # Test mapping creation - user_mapping, item_mapping = get_global_user_item_mapping("ml-1m", "/fake/path") - + user_mapping, item_mapping = get_global_user_item_mapping( + "ml-1m", "/fake/path" + ) + assert len(user_mapping) == 3 # Users 1, 2, 3 assert len(item_mapping) == 3 # Items 10, 20, 30 - - # Test time range computation + + # Test time range computation start_time, end_time = get_start_end_time("ml-1m", "/fake/path") - + assert start_time == 1000 assert end_time == 4000 - + # Test next day computation next_day = to_next_day(start_time) - assert next_day == start_time + 86400 \ No newline at end of file + assert next_day == start_time + 86400 diff --git a/tests/unit/test_utils_nc.py b/tests/unit/test_utils_nc.py index 53a9cac..634f843 100644 --- a/tests/unit/test_utils_nc.py +++ b/tests/unit/test_utils_nc.py @@ -1,28 +1,29 @@ -import pytest -import torch +import os +import tempfile +from unittest.mock import MagicMock, Mock, patch + import numpy as np +import pytest import scipy.sparse as sp -from unittest.mock import Mock, patch, MagicMock -import tempfile -import os +import torch from fedgraph.utils_nc import ( - normalize, - intersect1d, - setdiff1d, - label_dirichlet_partition, community_partition_non_iid, - get_in_comm_indexes, get_1hop_feature_sum, + get_in_comm_indexes, increment_dir, + intersect1d, + label_dirichlet_partition, + normalize, save_all_trainers_data, save_trainer_data_to_hugging_face, + setdiff1d, ) class TestUtilsNCBasicFunctions: """Test basic utility functions in utils_nc.""" - + def test_normalize_sparse_matrix(self): """Test normalize function for sparse matrices.""" # Create a simple sparse matrix @@ -30,14 +31,14 @@ def test_normalize_sparse_matrix(self): row = np.array([0, 0, 1, 1]) col = np.array([0, 1, 0, 1]) matrix = sp.csc_matrix((data, (row, col)), shape=(2, 2)) - + normalized = normalize(matrix) - + assert isinstance(normalized, sp.csr_matrix) # Check that rows sum to 1 (within tolerance) row_sums = np.array(normalized.sum(axis=1)).flatten() assert np.allclose(row_sums, 1.0) - + def test_normalize_with_zero_rows(self): """Test normalize function with zero rows.""" # Create matrix with a zero row @@ -45,68 +46,68 @@ def test_normalize_with_zero_rows(self): row = np.array([0, 0]) col = np.array([0, 1]) matrix = sp.csc_matrix((data, (row, col)), shape=(2, 2)) - + normalized = normalize(matrix) - + assert isinstance(normalized, sp.csr_matrix) # First row should sum to 1, second row should remain 0 row_sums = np.array(normalized.sum(axis=1)).flatten() assert np.allclose(row_sums[0], 1.0) assert np.allclose(row_sums[1], 0.0) - + def test_intersect1d_basic(self): """Test intersect1d function with basic tensors.""" t1 = torch.tensor([1, 2, 3, 4]) t2 = torch.tensor([3, 4, 5, 6]) - + result = intersect1d(t1, t2) - + assert torch.equal(result.sort().values, torch.tensor([3, 4])) - + def test_intersect1d_no_intersection(self): """Test intersect1d with no common elements.""" t1 = torch.tensor([1, 2]) t2 = torch.tensor([3, 4]) - + result = intersect1d(t1, t2) - + assert len(result) == 0 - + def test_intersect1d_all_same(self): """Test intersect1d with identical tensors.""" t1 = torch.tensor([1, 2, 3]) t2 = torch.tensor([1, 2, 3]) - + result = intersect1d(t1, t2) - + assert torch.equal(result.sort().values, torch.tensor([1, 2, 3])) def test_setdiff1d_basic(self): """Test setdiff1d function with basic tensors.""" t1 = torch.tensor([1, 2, 3, 4]) t2 = torch.tensor([3, 4, 5, 6]) - + result = setdiff1d(t1, t2) - + expected = torch.tensor([1, 2, 5, 6]) assert torch.equal(result.sort().values, expected.sort().values) - + def test_setdiff1d_no_difference(self): """Test setdiff1d with identical tensors.""" t1 = torch.tensor([1, 2, 3]) t2 = torch.tensor([1, 2, 3]) - + result = setdiff1d(t1, t2) - + assert len(result) == 0 - + def test_setdiff1d_complete_difference(self): """Test setdiff1d with no common elements.""" t1 = torch.tensor([1, 2]) t2 = torch.tensor([3, 4]) - + result = setdiff1d(t1, t2) - + expected = torch.tensor([1, 2, 3, 4]) assert torch.equal(result.sort().values, expected.sort().values) @@ -177,7 +178,7 @@ def test_bulk_save_passes_consistent_global_metadata(self, mock_save_trainer): class TestLabelDirichletPartition: """Test label_dirichlet_partition function.""" - + def test_label_dirichlet_partition_basic(self): """Test basic label Dirichlet partition.""" # Create simple labels @@ -186,22 +187,22 @@ def test_label_dirichlet_partition_basic(self): K = 3 # number of classes n_parties = 3 beta = 0.5 - + result = label_dirichlet_partition(labels, N, K, n_parties, beta) - + assert isinstance(result, list) assert len(result) == n_parties - + # Check that all indices are covered all_indices = set() for party_indices in result: all_indices.update(party_indices) assert len(all_indices) == N - + # Check that each party has some data for party_indices in result: assert len(party_indices) > 0 - + def test_label_dirichlet_partition_distribution_types(self): """Test different distribution types.""" labels = np.array([0, 1, 2] * 20) # 60 samples @@ -209,39 +210,39 @@ def test_label_dirichlet_partition_distribution_types(self): K = 3 n_parties = 3 beta = 1.0 - + # Test average distribution result_avg = label_dirichlet_partition( labels, N, K, n_parties, beta, distribution_type="average" ) - + # Test degree distribution result_degree = label_dirichlet_partition( labels, N, K, n_parties, beta, distribution_type="degree" ) - + assert len(result_avg) == n_parties assert len(result_degree) == n_parties - + # Results should be different (with high probability) assert result_avg != result_degree - + def test_label_dirichlet_partition_beta_values(self): """Test different beta values.""" labels = np.array([0, 1, 2] * 20) N = len(labels) K = 3 n_parties = 3 - + # Low beta (more heterogeneous) result_low = label_dirichlet_partition(labels, N, K, n_parties, beta=0.1) - + # High beta (more homogeneous) result_high = label_dirichlet_partition(labels, N, K, n_parties, beta=10.0) - + assert len(result_low) == n_parties assert len(result_high) == n_parties - + # Both should cover all indices assert sum(len(party) for party in result_low) == N assert sum(len(party) for party in result_high) == N @@ -249,7 +250,7 @@ def test_label_dirichlet_partition_beta_values(self): class TestCommunityPartition: """Test community_partition_non_iid function.""" - + def test_community_partition_non_iid(self): """Test community-based partition.""" # Create test data @@ -258,22 +259,22 @@ def test_community_partition_non_iid(self): num_clients = 3 nclass = 3 args_cuda = False - + result = community_partition_non_iid( non_iid_percent, labels, num_clients, nclass, args_cuda ) - + assert isinstance(result, list) assert len(result) == num_clients - + # Check that all indices are covered all_indices = set() for client_indices in result: all_indices.update(client_indices) - + # Should cover most or all of the data points assert len(all_indices) <= len(labels) - + # Each client should have some data (though this may not always be true) non_empty_clients = sum(1 for client in result if len(client) > 0) assert non_empty_clients >= 1 @@ -281,39 +282,41 @@ def test_community_partition_non_iid(self): class TestGetInCommIndexes: """Test get_in_comm_indexes function.""" - + def test_get_in_comm_indexes_basic(self): """Test basic functionality of get_in_comm_indexes.""" # Create simple graph - edge_index = torch.tensor([[0, 1, 2, 3, 4], - [1, 2, 3, 4, 0]]) - + edge_index = torch.tensor([[0, 1, 2, 3, 4], [1, 2, 3, 4, 0]]) + # Create node splits for trainers split_node_indexes = [ torch.tensor([0, 1]), torch.tensor([2, 3]), - torch.tensor([4]) + torch.tensor([4]), ] - + n_trainers = 3 num_hops = 2 idx_train = torch.tensor([0, 2, 4]) idx_test = torch.tensor([1, 3]) - + result = get_in_comm_indexes( - edge_index, split_node_indexes, n_trainers, - num_hops, idx_train, idx_test + edge_index, split_node_indexes, n_trainers, num_hops, idx_train, idx_test ) - + assert len(result) == 4 # Should return 4 elements - (communicate_node_global_indexes, in_com_train_node_local_indexes, - in_com_test_node_local_indexes, global_edge_indexes_clients) = result - + ( + communicate_node_global_indexes, + in_com_train_node_local_indexes, + in_com_test_node_local_indexes, + global_edge_indexes_clients, + ) = result + assert isinstance(communicate_node_global_indexes, list) assert isinstance(in_com_train_node_local_indexes, list) assert isinstance(in_com_test_node_local_indexes, list) assert isinstance(global_edge_indexes_clients, list) - + assert len(communicate_node_global_indexes) == n_trainers assert len(global_edge_indexes_clients) == n_trainers @@ -337,94 +340,93 @@ def test_get_in_comm_indexes_rejects_unsupported_one_hop(self): class TestGet1hopFeatureSum: """Test get_1hop_feature_sum function.""" - + def test_get_1hop_feature_sum_basic(self): """Test basic 1-hop feature sum calculation.""" # Create test data features = torch.randn(5, 10) # 5 nodes, 10 features - edge_index = torch.tensor([[0, 1, 2, 3], - [1, 2, 3, 4]]) + edge_index = torch.tensor([[0, 1, 2, 3], [1, 2, 3, 4]]) local_node_index = torch.tensor([0, 1, 2]) communicate_node_index = torch.tensor([0, 1, 2, 3, 4]) - + result = get_1hop_feature_sum( features, edge_index, local_node_index, communicate_node_index ) - + assert isinstance(result, torch.Tensor) assert result.shape[0] == len(communicate_node_index) assert result.shape[1] == features.shape[1] - + def test_get_1hop_feature_sum_isolated_nodes(self): """Test 1-hop feature sum with isolated nodes.""" features = torch.randn(3, 5) edge_index = torch.tensor([[0], [1]]) # Only one edge local_node_index = torch.tensor([0, 1, 2]) communicate_node_index = torch.tensor([0, 1, 2]) - + result = get_1hop_feature_sum( features, edge_index, local_node_index, communicate_node_index ) - + assert isinstance(result, torch.Tensor) assert result.shape == (3, 5) class TestIncrementDir: """Test increment_dir function.""" - + def test_increment_dir_new_directory(self, temp_dir): """Test increment_dir with new directory.""" base_path = os.path.join(temp_dir, "test_exp") - + result = increment_dir(base_path) - + assert result == base_path - + def test_increment_dir_existing_directory(self, temp_dir): """Test increment_dir with existing directories.""" base_path = os.path.join(temp_dir, "test_exp") - + # Create existing directories os.makedirs(base_path) os.makedirs(f"{base_path}1") - + result = increment_dir(base_path) - + assert result == f"{base_path}2" - + def test_increment_dir_with_comment(self, temp_dir): """Test increment_dir with comment.""" base_path = os.path.join(temp_dir, "test_exp") comment = "test_comment" - + result = increment_dir(base_path, comment) - + assert comment in result assert base_path in result class TestUtilsNCIntegration: """Integration tests for utils_nc functions.""" - + def test_tensor_operations_consistency(self): """Test consistency between intersect1d and setdiff1d.""" t1 = torch.tensor([1, 2, 3, 4, 5]) t2 = torch.tensor([3, 4, 5, 6, 7]) - + intersection = intersect1d(t1, t2) difference = setdiff1d(t1, t2) - + # Union of intersection and difference should cover both sets all_elements = torch.cat([intersection, difference]) original_elements = torch.cat([t1, t2]) - + # Check that all unique elements are covered unique_all = all_elements.unique() unique_original = original_elements.unique() - + assert len(unique_all) == len(unique_original) - + def test_partition_completeness(self): """Test that partition functions cover all data points.""" # Create test labels @@ -433,60 +435,60 @@ def test_partition_completeness(self): K = 5 n_parties = 4 beta = 1.0 - + partition = label_dirichlet_partition(labels, N, K, n_parties, beta) - + # Check completeness all_indices = set() for party in partition: all_indices.update(party) - + assert len(all_indices) == N assert all_indices == set(range(N)) - + # Check no overlaps total_assignments = sum(len(party) for party in partition) assert total_assignments == N - + def test_graph_communication_workflow(self): """Test complete workflow for graph communication setup.""" # Create a small graph num_nodes = 10 - edge_index = torch.tensor([[i, (i+1) % num_nodes] for i in range(num_nodes)]).t() - + edge_index = torch.tensor( + [[i, (i + 1) % num_nodes] for i in range(num_nodes)] + ).t() + # Create splits split_size = num_nodes // 3 split_node_indexes = [ - torch.arange(i*split_size, min((i+1)*split_size, num_nodes)) + torch.arange(i * split_size, min((i + 1) * split_size, num_nodes)) for i in range(3) ] - + # Handle remainder if sum(len(split) for split in split_node_indexes) < num_nodes: - split_node_indexes[-1] = torch.cat([ - split_node_indexes[-1], - torch.arange(split_size * 3, num_nodes) - ]) - + split_node_indexes[-1] = torch.cat( + [split_node_indexes[-1], torch.arange(split_size * 3, num_nodes)] + ) + n_trainers = 3 num_hops = 2 idx_train = torch.arange(0, 7) idx_test = torch.arange(7, 10) - + # Test the communication index computation result = get_in_comm_indexes( - edge_index, split_node_indexes, n_trainers, - num_hops, idx_train, idx_test + edge_index, split_node_indexes, n_trainers, num_hops, idx_train, idx_test ) - - (communicate_indexes, train_indexes, test_indexes, edge_indexes) = result - + + communicate_indexes, train_indexes, test_indexes, edge_indexes = result + # Verify structure assert len(communicate_indexes) == n_trainers assert len(train_indexes) == n_trainers assert len(test_indexes) == n_trainers assert len(edge_indexes) == n_trainers - + # Verify all trainers have some data for i in range(n_trainers): assert i in communicate_indexes