Make example scripts self-consistent in naming and distributed behavior

[kmontemayor2-sc]

Summary

A consistency pass over the 11 example train/infer scripts (plus tests/api_test), making every script say and do the same thing the same way, followed by adopting the new GbmlConfigPbWrapper typed feature maps (#643). Changes fall into six buckets:

1. Renames — one canonical name per concept

Concept	Canonical	Replaced
Processes per machine (config key / field / CLI)	local_world_size	num_inference_processes_per_machine, KDD --process_count
Spawned process index	local_rank	process_number (KDD)
Sampling workers (config key / field)	sampling_workers_per_process	sampling_workers_per_inference_process
Model artifact field	model_uri	model_state_dict_uri, saved_model_uri / saved_model_path
Output shard filename	machine_{...}_local_process_{local_rank}	..._local_process_number_{...}

Also fixed the spillter_kwargs → splitter_kwargs docstring typo. The KDD_2025 tutorial got renames only — its deliberately simple logic is untouched.

2. local_world_size resolution — one rule everywhere

• Colocated scripts share one identical block: explicit config value wins → torch.cuda.device_count() when CUDA → else a per-file DEFAULT_CPU_BASED_LOCAL_WORLD_SIZE constant (4 for inference, 1 for training).
• Always raise ValueError on GPU oversubscription — the homogeneous inference scripts previously only warned.
• Graph-store scripts now take local_world_size directly from cluster_info.num_processes_per_compute (the launcher-set COMPUTE_CLUSTER_LOCAL_WORLD_SIZE) instead of a config arg — the old CPU default of 2 silently disagreed with the launcher's 1, desyncing spawned processes from the compute
  process group.
• ⚠️ Behavior note: colocated LP training previously defaulted to 1 process even on GPU machines; it now defaults to GPU count, so e2e loss curves / step counts will shift on multi-GPU machines.

3. Distributed init — one canonical sequence

Every colocated spawned process now does: get device → cuda.set_device → init_process_group (rank/world_size computed inline via the machine formula, the only place it appears) → rank = torch.distributed.get_rank() / world_size = get_world_size() locals used everywhere downstream. Lead-process
guards are uniformly if rank == 0:. Graph-store heterogeneous inference now also reads rank from the process group after init_compute_process instead of a manual formula (its redundant machine_rank/machine_world_size fields are gone).

4. Cleanup — one canonical teardown

• Colocated inference: barrier → loader.shutdown() → gc.collect() → destroy_process_group() (LP files were missing the destroy).
• Graph-store: loader.shutdown() → shutdown_compute_process() → gc.collect() (GS homogeneous inference was missing shutdown_compute_process).
• GS homogeneous inference's BigQuery load moved out of the spawned process into orchestration, guarded by cluster_info.compute_node_rank == 0, matching its hetero sibling.
• LP heterogeneous inference orchestration gained the missing temporary-gloo-group destroy.
• All five non-KDD inference orchestrations now call mp.set_start_method("spawn") like the training scripts always have.

5. Dead/deprecated API usage removed

Dropped the deprecated local_process_rank / local_process_world_size loader kwargs from every example call site — they are silently ignored unless a (also-deprecated) context is passed, so they only taught users a dead parameter.

6. Adopt GbmlConfigPbWrapper typed feature maps (#643)

All hand-built {node,edge}_type_to_feature_dim dict comprehensions and condensed-key lookups (8 scripts, 2 api_test files, cora.ipynb, dblp.ipynb) replaced with gbml_config_pb_wrapper.node_type_to_feature_dim_map / .edge_type_to_feature_dim_map.

Investigated, not a bug

graph_store/heterogeneous_training.py derives query_node_type = supervision_edge_type.dst_node_type — reversed vs. standard mode. This is deliberate compensation for the reversed storage-side splitter_kwargs.supervision_edge_types in the GS DBLP config (traced through
DistNodeAnchorLinkSplitter's direction handling), and is now documented with comments instead of "fixed".

Verification

• make type_check and make format_py clean; repo-wide greps show zero stale occurrences of any old name.
• make unit_test_py PY_TEST_FILES="config_validator_test.py" (imports the example configs): 4/4 pass.
• Dual final review — an adversarial whole-branch review (MERGE-READY) and a Codex review; every confirmed finding was fixed on-branch.

Known follow-ups (out of scope)

• Optional hardening: guard mp.set_start_method against pre-initialized contexts (equally applies to the pre-existing training pattern).
• Library-level: finish deprecating/removing the loader local_process_rank/local_process_world_size params.
• Pre-existing on main: a dblp.ipynb cell references node_feature_dims[node_type] where node_type may be undefined.

[kmontemayor2-sc]

/all_test

[github-actions]

GiGL Automation

@ 18:33:43UTC : 🔄 Scala Unit Test started.

@ 18:42:42UTC : ✅ Workflow completed successfully.

[github-actions]

GiGL Automation

@ 18:33:43UTC : 🔄 Python Unit Test started.

@ 19:42:58UTC : ✅ Workflow completed successfully.

[github-actions]

GiGL Automation

@ 18:33:46UTC : 🔄 Integration Test started.

@ 20:25:50UTC : ✅ Workflow completed successfully.

[github-actions]

GiGL Automation

@ 18:33:47UTC : 🔄 E2E Test started.

@ 20:09:58UTC : ✅ Workflow completed successfully.

[github-actions]

GiGL Automation

@ 18:33:47UTC : 🔄 Lint Test started.

@ 18:41:46UTC : ✅ Workflow completed successfully.

[github-actions]

GiGL Automation

@ 18:33:48UTC : 🔄 C++ Unit Test started.

@ 18:35:56UTC : ✅ Workflow completed successfully.