diff --git a/docs/plans/20260730-graph-store-multiple-supervision-edge-types.md b/docs/plans/20260730-graph-store-multiple-supervision-edge-types.md new file mode 100644 index 000000000..224e72325 --- /dev/null +++ b/docs/plans/20260730-graph-store-multiple-supervision-edge-types.md @@ -0,0 +1,443 @@ +# GraphStore multiple-supervision-edge support + +## Objective + +Enable one GraphStore `DistABLPLoader` to fetch, sample, and return labels for multiple heterogeneous supervision edge +types that share one anchor node type, without regressing the existing one-edge, labeled-homogeneous, sharded-fetch, or +incoming-edge behavior. + +The task metadata and graph materialization layers are already plural: the task proto stores repeated supervision edge +types, its wrapper returns a list, and the dataset factory passes that list into the splitter +(`proto/snapchat/research/gbml/gbml_config.proto:18-34`, `gigl/src/common/types/pb_wrappers/task_metadata.py:58-88`, +`gigl/distributed/dataset_factory.py:568-585`). The singular RPC/fetch boundary is the main TODO +(`gigl/distributed/graph_store/remote_dist_dataset.py:380-446`), but correct per-type negative output also requires a +loader change because positive and negative results are currently flattened independently +(`gigl/distributed/dist_ablp_neighborloader.py:876-886`). + +## Scope and compatibility decisions + +### In scope + +1. Widen `RemoteDistDataset.fetch_ablp_input()` from one supervision edge type to one or many, while keeping scalar + input source-compatible (`gigl/distributed/graph_store/remote_dist_dataset.py:447-586`). +2. Batch all requested edge types into the existing one-request-per-storage-server fan-out, rather than issuing one RPC + per edge type (`gigl/distributed/graph_store/remote_dist_dataset.py:388-426`). +3. Return the existing `ABLPInputNodes` schema, whose `labels` member is already a dictionary from supervision edge type + to `(positive, optional negative)` tensors (`gigl/utils/sampling.py:91-142`). +4. Remove the GraphStore-only one-edge loader restriction and validate the full per-server schema before constructing + sampler inputs (`gigl/distributed/dist_ablp_neighborloader.py:676-713`). +5. Support per-edge negative-label availability. Today the loader compares the whole `(positive, negative)` tuple with + `None`, so any label entry makes the global `has_negatives` flag true + (`gigl/distributed/dist_ablp_neighborloader.py:684-703`, `gigl/utils/sampling.py:139-142`). +6. Prove both the RPC data path and final multi-key `y_positive`/`y_negative` output, including `edge_dir="in"`, + sharding, and an edge type without hard negatives. +7. Align the preferred public API with colocated mode's anchor-outward supervision types while retaining the existing + incoming GraphStore scalar call as a compatibility form. + +### Direction contract + +The preferred public contract will match colocated mode: caller-provided supervision types are anchor-outward, so every +type has `src_node_type == anchor_node_type`. Colocated mode already enforces that rule and reverses the types +internally for incoming sampling (`gigl/distributed/dist_ablp_neighborloader.py:471-490`). Canonical dataset +construction also starts from the task-metadata list and reverses label edges before registering an incoming graph +(`gigl/distributed/dataset_factory.py:568-585`, `gigl/types/graph.py:110-131`). + +For `edge_dir="in"`, `RemoteDistDataset` will therefore translate the canonical outward list to registered sampling +orientation before constructing the internal request: + +```text +public/task type registered request key final loader key +A --r1--> B B --r1--> A A --r1--> B +A --r2--> C C --r2--> A A --r2--> C +``` + +The final reversal already occurs during collation (`gigl/distributed/dist_ablp_neighborloader.py:905-921`). For +`edge_dir="out"`, the public, registered, and final keys are identical. + +Backward compatibility is necessary because the current GraphStore example treats the destination of `paper -> author` +as the anchor and passes that stored orientation directly +(`examples/link_prediction/graph_store/heterogeneous_training.py:180-202`); its storage config documents a manual +reverse of the splitter input +(`examples/link_prediction/graph_store/configs/e2e_het_dblp_sup_gs_task_config.yaml:55-67`). For incoming graphs only, +accept a legacy request when all supplied types have `dst_node_type == anchor_node_type`; those types are already +registered orientation and are not reversed again. Reject a list that mixes canonical and legacy orientations. When both +endpoints have the anchor type, reversal produces the same node-type endpoints, so resolve by positive-label topology +and reject only if neither candidate exists. This codifies an implicit convention as new validation; the current client +checks only optional-argument pairing (`gigl/distributed/graph_store/remote_dist_dataset.py:545-558`). + +### Explicit non-goals + +- No proto, partition format, or graph registration change. `EdgeType` is already hashable, graph labels are stored by + edge type, and partitioned label maps are already dictionaries (`gigl/src/common/types/graph_data.py:24-33`, + `gigl/types/graph.py:60-91`, `gigl/distributed/dist_partitioner.py:150-178`). +- No change to neighbor/PPR sampling algorithms. `ABLPNodeSamplerInput` already slices and shares every dictionary entry + (`gigl/distributed/sampler.py:13-75`), and the base sampler already visits all positive and negative label types when + building seeds and metadata (`gigl/distributed/base_sampler.py:129-199`). +- No promise that the example trainer can optimize multiple objectives. It still rejects task metadata containing more + than one type and assumes `main_data.y_positive` is a single tensor + (`examples/link_prediction/graph_store/heterogeneous_training.py:927-934`, `:297-315`). Updating loss aggregation, + random-negative semantics, and trainer configuration should be a follow-up after the core loader contract is proven. +- No mixed-version compute/storage rollout. GraphStore RPC sends a Python callable and Python arguments directly, so the + request/response dataclass change assumes compute and storage use the same release + (`gigl/distributed/graph_store/compute.py:99-124`). + +## Current data flow and the actual choke points + +1. `RemoteDistDataset._fetch_ablp_input()` builds one request per selected server, but each request contains one + `supervision_edge_type` and each future returns one `(anchors, positive, negative)` tuple + (`gigl/distributed/graph_store/remote_dist_dataset.py:380-444`). +2. `FetchABLPInputRequest` encodes that singular field (`gigl/distributed/graph_store/messages.py:90-117`). +3. `DistServer.get_ablp_input()` fetches anchors once, resolves one positive and optional negative label topology, and + returns one label pair (`gigl/distributed/graph_store/dist_server.py:542-576`). +4. The public fetch method wraps that one pair into an `ABLPInputNodes.labels` dictionary, despite the dictionary + already being able to carry several keys (`gigl/distributed/graph_store/remote_dist_dataset.py:545-586`, + `gigl/utils/sampling.py:139-142`). +5. GraphStore loader setup reads schema only from the first server, computes a faulty global negative flag, then + explicitly rejects any key count other than one (`gigl/distributed/dist_ablp_neighborloader.py:676-713`). +6. Everything after that guard is already structured as dictionaries: loader conversion loops over all label entries + (`gigl/distributed/dist_ablp_neighborloader.py:718-762`), sampling adds every label type to metadata + (`gigl/distributed/base_sampler.py:152-199`), remapping handles dictionaries + (`gigl/distributed/utils/ablp.py:97-160`), and output is flattened only for one result while multiple positive + results remain dictionaries (`gigl/distributed/dist_ablp_neighborloader.py:838-886`). +7. Positive and negative outputs are flattened independently. With two positive types but hard negatives for only one + type, `y_negative` becomes unkeyed and loses its relation (`gigl/distributed/dist_ablp_neighborloader.py:876-886`). + +The colocated path demonstrates the intended cardinality behavior. It normalizes a scalar or non-empty list +(`gigl/distributed/dist_ablp_neighborloader.py:263-274`), resolves and fetches labels once per type +(`gigl/distributed/dist_ablp_neighborloader.py:538-570`), and already has positive-only, positive-and-negative, +same-endpoint/different-relation, and incoming-direction multi-type cases +(`tests/unit/distributed/dist_ablp_neighborloader_test.py:1095-1277`). + +## Proposed implementation + +### 1. Make the RPC contract plural and reuse `ABLPInputNodes` + +Files: + +- `gigl/distributed/graph_store/messages.py` +- `gigl/distributed/graph_store/dist_server.py` +- `tests/unit/distributed/graph_store/messages_test.py` +- `tests/unit/distributed/dist_server_test.py` + +Change `FetchABLPInputRequest.supervision_edge_type` to `supervision_edge_types: tuple[EdgeType, ...]`. A tuple keeps +the frozen request structurally immutable. The field represents registered sampling-orientation types; public +orientation resolution happens before request construction. Update its documentation/examples. + +Change `DistServer.get_ablp_input()` to return `ABLPInputNodes`: + +1. Before fetching anchors, reject an empty tuple, duplicates, and a registered type whose effective anchor endpoint is + wrong (`src` for `out`, `dst` for `in`). This repeats validation at the authoritative RPC boundary because callers + can invoke the server helper without the public normalizer (`gigl/distributed/graph_store/compute.py:99-124`). +2. Fetch and slice anchors exactly once. +3. Iterate `request.supervision_edge_types` in caller order. +4. For each type, call `select_label_edge_types()` against the registered edge types, then + `get_labels_for_anchor_nodes()` with the same anchors and `max_labels_per_anchor_node`. +5. Store the returned pair under the registered supervision edge type in `labels`. +6. Return one `ABLPInputNodes(anchor_node_type=request.node_type, ...)`. + +`get_labels_for_anchor_nodes()` already defines the required per-type contract: positive and optional negative tensors +are `[N, M]`, padding uses `-1`, and an empty anchor set returns `[0, 0]` tensors while preserving whether a negative +topology exists (`gigl/utils/data_splitters.py:607-682`). + +This preserves one RPC per selected server. The client currently dispatches one future per server +(`gigl/distributed/graph_store/remote_dist_dataset.py:415-426`); an edge-by-server fan-out would multiply calls while +the client RPC pool defaults to four workers (`gigl/distributed/graph_store/compute.py:127-175`). + +### 2. Normalize and validate the public fetch API + +Files: + +- `gigl/distributed/graph_store/remote_dist_dataset.py` +- `gigl/distributed/graph_store/dist_server.py` + +Widen the existing keyword without renaming it: + +```python +supervision_edge_type: Optional[Union[EdgeType, list[EdgeType]]] = None +``` + +Normalize `None` to `[DEFAULT_HOMOGENEOUS_EDGE_TYPE]`, a scalar to a one-item list, and a list to a copied list. Reject: + +- an empty list; +- duplicate edge types; +- only one of `anchor_node_type` and `supervision_edge_type` being supplied; +- a heterogeneous list that is neither uniformly canonical anchor-outward nor uniformly legacy incoming registered + orientation; +- a list mixing canonical and legacy incoming orientations; +- more than one type for the labeled-homogeneous default. + +Fetch `edge_dir` and registered edge types once. Resolve canonical outward inputs to registered orientation, verify +positive-label topology for every resolved type, and record per-type negative-topology presence with +`select_label_edge_types()` (`gigl/types/graph.py:349-372`). Preserve a uniformly legacy incoming list as described +above. Pass only the resolved tuple into every per-server request, and have `_fetch_ablp_input()` return +`dict[int, ABLPInputNodes]` directly. + +Keep the public fetch result dense for compatibility, even though the loader can already accept an in-range sparse +mapping and synthesize missing-rank inputs (`gigl/distributed/dist_ablp_neighborloader.py:663-675`, `:718-762`). For +every unrequested rank, construct a topology-complete placeholder: empty anchors, one `[0, 0]` positive tensor per +resolved type, and either a `[0, 0]` negative tensor or `None` according to that type's registered negative topology. +This removes the need to infer placeholder provenance later; a provided `None` where topology requires negatives is +uniformly invalid. + +Retain request order in all dictionaries for deterministic logs and tests, while using set equality only for schema +validation. + +### 3. Make GraphStore loader setup schema-safe and per-edge + +File: `gigl/distributed/dist_ablp_neighborloader.py` + +Replace the first-entry/global-boolean logic with full validation: + +1. Require a non-empty `input_nodes` mapping and non-empty label schema. +2. Establish the expected `anchor_node_type` and ordered supervision-key tuple from the first entry. +3. For every server entry, require the same anchor type, the same supervision key set, a one-dimensional integral anchor + tensor, and two-dimensional integral positive/negative tensors whose first dimension equals the anchor count. + Remapping indexes dimension one, so accepting a one-dimensional label tensor would only defer failure + (`gigl/distributed/utils/ablp.py:62-76`). +4. Validate every supervision type has the effective anchor endpoint required by the dataset's `edge_dir`. +5. Resolve each positive and optional negative label edge type independently against `edge_types` using + `select_label_edge_types()`. This makes topology, not “some server returned a tensor,” the source of truth. +6. Uniformly require a positive tensor and require a negative tensor exactly when that edge type has registered negative + topology; remote placeholders are already schema-complete. +7. Remove the heterogeneous GraphStore `len(...) != 1` guard at `gigl/distributed/dist_ablp_neighborloader.py:706-713`. +8. Preserve the later labeled-homogeneous exactly-one invariant + (`gigl/distributed/dist_ablp_neighborloader.py:923-930`). + +This also fixes mixed negative availability: `_negative_label_edge_types` becomes the subset selected from graph +topology, and each `ABLPNodeSamplerInput` receives only its actual negative keys. No single `has_negatives` flag should +control all edge types. + +Prefer extracting the validation/normalization into a small private helper that returns a validated schema plus the +dense `list[ABLPNodeSamplerInput]`. That makes failure cases testable without starting sampling workers. + +Move local validation before GraphStore worker-option creation and port allocation. The current order fetches a port +before the input schema is fully checked (`gigl/distributed/dist_ablp_neighborloader.py:643-678`, +`gigl/distributed/base_dist_loader.py:600-624`). Gather a compact local success/error result across compute ranks; if +any rank failed validation, make all ranks raise before collective port allocation or backend registration. + +Continue accepting sparse in-range `input_nodes` mappings. Only negative or out-of-range server ranks are rank errors; +missing ranks are not. + +### 4. Preserve relation keys for mixed negative output + +File: `gigl/distributed/dist_ablp_neighborloader.py` + +Change `_set_labels()` flattening to depend on supervision cardinality: + +- if the loader has one supervision type, retain the existing bare tensor/ragged-dictionary `y_positive` and + `y_negative` API; +- if it has more than one supervision type, keep both outputs edge-keyed; `y_negative` contains only the types with + negative topology, even when that subset has length one; +- if no type has negatives, continue omitting `y_negative`. + +This is required because the current independent `len(output_negative_labels)` check loses the only negative edge key in +the mixed case (`gigl/distributed/dist_ablp_neighborloader.py:876-886`). + +### 5. Update API documentation and logging + +Files: + +- `gigl/distributed/graph_store/remote_dist_dataset.py` +- `gigl/distributed/graph_store/messages.py` +- `gigl/distributed/graph_store/dist_server.py` +- `gigl/distributed/dist_ablp_neighborloader.py` + +Document: + +- scalar and list inputs; +- canonical anchor-outward input, legacy incoming compatibility, and rejection of mixed orientation; +- the exact public key -> registered key -> final key mapping for `edge_dir="in"`; +- dictionary output for multiple final label types; +- per-type optional negatives; +- labeled-homogeneous remaining single-type. + +Log the ordered list of requested supervision types and the independently resolved positive/negative label types. Do not +log label tensors or node IDs. + +## Test plan + +### Unit: request and server fetch contract + +Extend `tests/unit/distributed/graph_store/messages_test.py`, whose current ABLP case asserts the singular request +(`tests/unit/distributed/graph_store/messages_test.py:33-55`): + +- plural tuple construction, equality, frozen behavior, and optional slice; +- one-item tuple as the backward-compatible normalized form. + +Update every existing ABLP case in `tests/unit/distributed/dist_server_test.py`. That tracked module already pins the +singular request and tuple response contract (`tests/unit/distributed/dist_server_test.py:345-595`). Add: + +- anchors are fetched/sliced once and reused for two supervision types; +- exact positive and negative tensor values are keyed by both types; +- one type with negatives plus one without negatives; +- empty anchors retain every key and the correct `None` versus `[0, 0]` negative shape; +- empty tuple, duplicates, wrong registered anchor endpoint, and missing label topology raise with the offending type. + +### Unit: remote fetch API and sharding + +Extend `tests/unit/distributed/graph_store/remote_dist_dataset_test.py`. Its existing cases already cover singleton +fetches, split selection, label caps, defaults, mismatched optional arguments, and exact sharding requests +(`tests/unit/distributed/graph_store/remote_dist_dataset_test.py:440-590`, `:651-735`, `:1029-1277`). + +Add: + +- scalar input produces a one-item request tuple and unchanged output; +- canonical incoming types are reversed before request construction, while the current scalar legacy incoming form + remains unchanged; +- mixed canonical/legacy incoming lists fail before dispatch; +- two types produce one request per assigned server, not one request per type/server pair; +- exact key order and exact positive/negative values for every server; +- dense empty placeholders preserve both positive keys and topology-correct negative tensors; +- empty list, duplicates, inconsistent anchor endpoint, missing type, and homogeneous multi-type errors; +- `edge_dir="out"` keeps canonical keys; `edge_dir="in"` proves both canonical translation and legacy compatibility; +- fractional server slicing applies identical row selection to every type. + +### Unit: GraphStore loader conversion and output + +Extend `tests/unit/distributed/dist_ablp_neighborloader_test.py` with tests around the extracted pure GraphStore setup +helper: + +- consistent two-type inputs become one `ABLPNodeSamplerInput` per server; +- mixed negative availability produces the exact positive and negative edge-type sets; +- empty unassigned servers receive schema-correct empty tensors; +- mismatched anchor types, key sets, out-of-range ranks, dimensions, integer dtypes, row counts, and negative topology + fail before worker startup; +- sparse in-range server mappings remain valid; +- labeled-homogeneous still rejects multiple types. + +Run shared loader/collation cases for each direction and for both `use_edge_index_output` modes. Assert: + +- both final `y_positive` keys and exact values; +- hard negatives on both types; +- the mixed case where only one type has hard negatives remains a one-key `y_negative` dictionary rather than a bare + value; +- incoming registered keys are reversed back to the exact canonical outward keys. + +Also fix the colocated multi-edge helper so the negative assertion is inside a loop over every negative edge type. It +currently loops over positives but checks only the final `edge_type` for negatives +(`tests/unit/distributed/dist_ablp_neighborloader_test.py:347-366`), despite its fixture declaring negative expectations +for both types (`tests/unit/distributed/dist_ablp_neighborloader_test.py:1152-1193`). + +### Integration: real GraphStore RPC and sampler + +Extend `tests/integration/distributed/graph_store/graph_store_integration_test.py`. The current test checks singleton +remote input and loader startup (`tests/integration/distributed/graph_store/graph_store_integration_test.py:252-325`); +the heterogeneous loader test is skipped and exercises only a neighbor loader +(`tests/integration/distributed/graph_store/graph_store_integration_test.py:1023-1049`). + +Do not rely on the current mocked-asset model: it stores only one `sample_edge_type`, and frozen config generation emits +a singleton supervision list (`gigl/src/mocking/lib/mocked_dataset_resources.py:142-151`, +`gigl/src/mocking/dataset_asset_mocker.py:316-327`). Instead: + +1. Generalize the test-only `create_heterogeneous_dataset_for_ablp()` builder to accept per-edge-type positive/negative + maps while retaining its scalar convenience form (`tests/test_assets/distributed/test_dataset.py:236-300`). +2. Add a top-level, multiprocessing-picklable builder for a tiny two-type dataset with one shared anchor type. +3. Add an optional test-only dataset-builder field to `ServerProcessArgs`; after initializing the storage process group, + `_run_storage_main_process()` uses it instead of `build_storage_dataset()`. The current harness otherwise always + builds from a task-config URI (`tests/integration/distributed/graph_store/graph_store_integration_test.py:709-728`, + `:761-767`). + +Use the existing incoming GraphStore direction for the real-RPC acceptance path. Do not reuse the skipped heterogeneous +neighbor-loader case, whose Cloud Build failure is unresolved +(`tests/integration/distributed/graph_store/graph_store_integration_test.py:1023-1049`). Through real RPC: + +1. Fetch the same split with two canonical outward task types. +2. Assert each storage response has the exact anchors, exact two registered-key label schema, and exact padded + positive/negative rows. +3. Start `DistABLPLoader` with `use_edge_index_output=True`, consume at least one batch, and assert both canonical + outward output keys and their exact global positive/negative pairs, including negatives for only one type. +4. Exercise a sharded compute rank with an unassigned storage server. + +Keep the existing singleton and homogeneous integration cases as regression coverage. Do not treat the current +multi-type dataset-storage test alone as sufficient: its supervision types have different anchor types +(`tests/integration/distributed/distributed_dataset_test.py:155-191`), so they cannot feed one ABLP loader under the +shared-anchor contract. + +## Verification commands + +Run targeted tests first: + +```bash +.venv/bin/python -m unittest \ + tests.unit.distributed.graph_store.messages_test \ + tests.unit.distributed.graph_store.remote_dist_dataset_test \ + tests.unit.distributed.dist_server_test \ + tests.unit.distributed.dist_ablp_neighborloader_test +``` + +Then run repository gates and the GraphStore integration target: + +```bash +make type_check +make check_format +make integration_test PY_TEST_FILES="graph_store_integration_test.py" +``` + +Finally run the full unit suite through `make unit_test_py` once the current worktree's unrelated dev-utils import is +repaired or removed. + +## Rollout risks and mitigations + +| Risk | Mitigation | +| --------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------ | +| Incoming-edge orientation is interpreted differently by callers | Prefer canonical anchor-outward input, translate to registered orientation, retain the uniform legacy incoming form, reject mixed forms, and pin exact final keys. | +| A malformed or stale server response reaches sampling workers | Validate every rank's anchor type, key set, dimensions, integer dtype, and topology before collective worker setup; repeat request validation server-side. | +| A type without hard negatives inherits another type's negative schema | Derive negative availability independently with `select_label_edge_types()` and test mixed availability. | +| Mixed negative output loses its only relation key | Flatten only for a single-supervision loader; keep multi-supervision outputs keyed. | +| Sharded ranks lose per-server alignment | Keep remote fetch dense with topology-complete placeholders while preserving the loader's existing sparse-input support. | +| RPC count grows with the number of supervision types | Batch all types in the one existing request to each selected server. | +| Integration coverage depends on singleton mocked assets or a skipped heterogeneous test | Inject a tiny programmatic multi-type dataset through the test harness. | +| Example users infer that multi-objective training is now supported | Keep the example trainer limitation explicit and track it as a separate follow-up. | + +## Acceptance criteria + +1. The old incoming scalar GraphStore call remains accepted and yields the same one-key `ABLPInputNodes` and final + tensor-shaped `y_positive`/`y_negative` behavior. +2. A canonical anchor-outward list taken from task metadata is translated to registered incoming orientation and fetched + in one RPC per selected storage server; mixed canonical/legacy lists are rejected. +3. Every server response and every `ABLPNodeSamplerInput` carries both positive label types; negative label keys match + registered topology per type. +4. A real-RPC incoming GraphStore `DistABLPLoader` batch exposes both canonical outward output keys with correct global + label pairs. Direct server/client and shared collation tests cover `edge_dir="out"` as well. +5. Rank/world-size sharding preserves identical anchor-row selection across all requested label tensors; remote fetch + remains dense and manual sparse loader input remains valid. +6. Empty lists, duplicates, inconsistent anchor endpoints, inconsistent per-server schemas, malformed tensor + dimensions/dtypes, and labeled-homogeneous multi-type requests fail with actionable errors before collective + sampling-worker setup. +7. With multiple supervision types and negatives for only one, `y_positive` and `y_negative` remain edge-keyed; + singleton loaders retain the existing flattened API. +8. Targeted unit tests, type checking, formatting, and the GraphStore integration test pass. + +## Baseline verification before implementation + +- `tests.unit.distributed.graph_store.remote_dist_dataset_test` passes directly. +- `tests.unit.distributed.dist_ablp_neighborloader_test` passes directly: 21 tests ran in 1001.292 seconds with one + skip. This exercises the existing colocated multi-edge cases at + `tests/unit/distributed/dist_ablp_neighborloader_test.py:1095-1332`. +- The normal `make unit_test_py` entry point is currently blocked before tests by an unrelated untracked file, + `gigl/utils/dev/tb_smoke_main.py:24`, importing a missing `gigl.utils.tensorboard_writer`. This is a worktree baseline + issue, not part of the GraphStore plan. + +## Dual-review resolution + +The initial draft received two independent read-only reviews: + +- SOL: `.claude/tmp/codex-verify/20260730-graph-store-multi-edge-plan/review.md` +- Claude Fable: `.claude/tmp/claude-fable-review/20260730-graph-store-multi-edge-plan/review.md` + +Both reviewers found the tracked server-test omission and placeholder ambiguity +(`.claude/tmp/codex-verify/20260730-graph-store-multi-edge-plan/review.md:124-153`, +`.claude/tmp/claude-fable-review/20260730-graph-store-multi-edge-plan/review.md:125-168`). The revised plan now updates +`tests/unit/distributed/dist_server_test.py` unconditionally and creates topology-complete placeholders. + +SOL identified three blocking design gaps: canonical incoming task types were not accepted, mixed negatives lost their +output key, and the proposed integration fixture was not constructible with singleton mocked assets +(`.claude/tmp/codex-verify/20260730-graph-store-multi-edge-plan/review.md:64-123`). The revised plan adopts +canonical-outward translation with legacy compatibility, changes flattening semantics for multi-type loaders, and +specifies a concrete programmatic integration-dataset hook. + +The plan also adopts the reviewers' non-blocking hardening findings: strict dimension/dtype validation, server-side +request validation, sparse loader input compatibility, validation before collective initialization, exact incoming key +mapping, and a test for mixed negatives through final collation +(`.claude/tmp/codex-verify/20260730-graph-store-multi-edge-plan/review.md:155-237`, +`.claude/tmp/claude-fable-review/20260730-graph-store-multi-edge-plan/review.md:171-236`). diff --git a/examples/link_prediction/graph_store/heterogeneous_training.py b/examples/link_prediction/graph_store/heterogeneous_training.py index 2be34e608..2c772236b 100644 --- a/examples/link_prediction/graph_store/heterogeneous_training.py +++ b/examples/link_prediction/graph_store/heterogeneous_training.py @@ -211,6 +211,7 @@ def _setup_dataloaders( channel_size=sampling_worker_shared_channel_size, process_start_gap_seconds=process_start_gap_seconds, shuffle=shuffle, + use_edge_index_output=True, ) print(f"---Rank {rank} finished setting up main loader for split={split}") @@ -294,23 +295,16 @@ def _compute_loss( ) # Extracting local query, random negative, positive, hard_negative, and random_negative indices. - query_node_idx: torch.Tensor = torch.arange( - main_data[query_node_type].batch_size - ).to(device) + query_node_idx = torch.arange(main_data[query_node_type].batch_size, device=device) random_negative_batch_size = random_negative_data[labeled_node_type].batch_size - positive_idx: torch.Tensor = torch.cat(list(main_data.y_positive.values())).to( - device - ) - repeated_query_node_idx = query_node_idx.repeat_interleave( - torch.tensor([len(v) for v in main_data.y_positive.values()]).to(device) - ) + positive_label_edge_index: torch.Tensor = main_data.y_positive + repeated_query_node_idx = query_node_idx[positive_label_edge_index[0]] + positive_idx = positive_label_edge_index[1] if hasattr(main_data, "y_negative"): - hard_negative_idx: torch.Tensor = torch.cat( - list(main_data.y_negative.values()) - ).to(device) + hard_negative_idx: torch.Tensor = main_data.y_negative[1] else: - hard_negative_idx = torch.empty(0, dtype=torch.long).to(device) + hard_negative_idx = torch.empty(0, dtype=torch.long, device=device) # Use local IDs to get the corresponding embeddings in the tensors diff --git a/examples/link_prediction/graph_store/homogeneous_training.py b/examples/link_prediction/graph_store/homogeneous_training.py index 2d4c22788..25ece94a0 100644 --- a/examples/link_prediction/graph_store/homogeneous_training.py +++ b/examples/link_prediction/graph_store/homogeneous_training.py @@ -241,6 +241,7 @@ def _setup_dataloaders( channel_size=sampling_worker_shared_channel_size, process_start_gap_seconds=process_start_gap_seconds, shuffle=shuffle, + use_edge_index_output=True, ) logger.info(f"---Rank {rank} finished setting up main loader for split={split}") @@ -302,21 +303,16 @@ def _compute_loss( main_embeddings = model(data=main_data, device=device) random_negative_embeddings = model(data=random_negative_data, device=device) - query_node_idx: torch.Tensor = torch.arange(main_data.batch_size).to(device) + query_node_idx = torch.arange(main_data.batch_size, device=device) random_negative_batch_size = random_negative_data.batch_size - positive_idx: torch.Tensor = torch.cat(list(main_data.y_positive.values())).to( - device - ) - repeated_query_node_idx = query_node_idx.repeat_interleave( - torch.tensor([len(v) for v in main_data.y_positive.values()]).to(device) - ) + positive_label_edge_index: torch.Tensor = main_data.y_positive + repeated_query_node_idx = query_node_idx[positive_label_edge_index[0]] + positive_idx = positive_label_edge_index[1] if hasattr(main_data, "y_negative"): - hard_negative_idx: torch.Tensor = torch.cat( - list(main_data.y_negative.values()) - ).to(device) + hard_negative_idx: torch.Tensor = main_data.y_negative[1] else: - hard_negative_idx = torch.empty(0, dtype=torch.long).to(device) + hard_negative_idx = torch.empty(0, dtype=torch.long, device=device) # Use local IDs to get the corresponding embeddings in the tensors diff --git a/examples/link_prediction/heterogeneous_training.py b/examples/link_prediction/heterogeneous_training.py index 8ed672b7c..e41a19191 100644 --- a/examples/link_prediction/heterogeneous_training.py +++ b/examples/link_prediction/heterogeneous_training.py @@ -144,6 +144,7 @@ def _setup_dataloaders( # This is done so that each process on the current machine which initializes a `main_loader` doesn't compete for memory, causing potential OOM process_start_gap_seconds=process_start_gap_seconds, shuffle=shuffle, + use_edge_index_output=True, ) logger.info(f"---Rank {rank} finished setting up main loader") @@ -218,25 +219,16 @@ def _compute_loss( # Local in this case refers to the local index in the batch, while global subsequently refers to the node's unique global ID across all nodes in the dataset. # Global ids are stored in data.node, ex. `data.node = [50, 20, 10]` where the `0` is the local index for global id `50`. Note that each global id in data.node is unique. # TODO (mkolodner-sc): If GLT migrates towards using PyG's `.n_id` field instead of `.node`, update this code. - query_node_idx: torch.Tensor = torch.arange( - main_data[query_node_type].batch_size - ).to(device) + query_node_idx = torch.arange(main_data[query_node_type].batch_size, device=device) random_negative_batch_size = random_negative_data[labeled_node_type].batch_size - # main_data.y_positive is a dict[query_node_local_index: int, labeled_node_local_indices: torch.Tensor], even in the heterogeneous setting. - positive_idx: torch.Tensor = torch.cat(list(main_data.y_positive.values())).to( - device - ) - # We also extract a repeated query node index tensor which upsamples each query node based on the number of positives it has - repeated_query_node_idx = query_node_idx.repeat_interleave( - torch.tensor([len(v) for v in main_data.y_positive.values()]).to(device) - ) + positive_label_edge_index: torch.Tensor = main_data.y_positive + repeated_query_node_idx = query_node_idx[positive_label_edge_index[0]] + positive_idx = positive_label_edge_index[1] if hasattr(main_data, "y_negative"): - hard_negative_idx: torch.Tensor = torch.cat( - list(main_data.y_negative.values()) - ).to(device) + hard_negative_idx: torch.Tensor = main_data.y_negative[1] else: - hard_negative_idx = torch.empty(0, dtype=torch.long).to(device) + hard_negative_idx = torch.empty(0, dtype=torch.long, device=device) # Use local IDs to get the corresponding embeddings in the tensors diff --git a/examples/link_prediction/homogeneous_training.py b/examples/link_prediction/homogeneous_training.py index b95a77489..a84b55114 100644 --- a/examples/link_prediction/homogeneous_training.py +++ b/examples/link_prediction/homogeneous_training.py @@ -134,6 +134,7 @@ def _setup_dataloaders( # This is done so that each process on the current machine which initializes a `main_loader` doesn't compete for memory, causing potential OOM process_start_gap_seconds=process_start_gap_seconds, shuffle=shuffle, + use_edge_index_output=True, ) logger.info(f"---Rank {rank} finished setting up main loader") @@ -187,23 +188,17 @@ def _compute_loss( # Extracting local query, random negative, positive, hard_negative, and random_negative indices. # Local in this case refers to the local index in the batch, while global subsequently refers to the node's unique global ID across all nodes in the dataset. # Global ids are stored in data.node, ex. `data.node = [50, 20, 10]` where the `0` is the local index for global id `50`. Note that each global id in data.node is unique. - query_node_idx: torch.Tensor = torch.arange(main_data.batch_size).to(device) + query_node_idx = torch.arange(main_data.batch_size, device=device) random_negative_batch_size = random_negative_data.batch_size - # main_data.y_positive is a dict[query_node_local_index: int, labeled_node_local_indices: torch.Tensor] - positive_idx: torch.Tensor = torch.cat(list(main_data.y_positive.values())).to( - device - ) - # We also extract a repeated query node index tensor which upsamples each query node based on the number of positives it has - repeated_query_node_idx = query_node_idx.repeat_interleave( - torch.tensor([len(v) for v in main_data.y_positive.values()]).to(device) - ) + # Label edge indices are [anchor_local_id, label_local_id] pairs. + positive_label_edge_index: torch.Tensor = main_data.y_positive + repeated_query_node_idx = query_node_idx[positive_label_edge_index[0]] + positive_idx = positive_label_edge_index[1] if hasattr(main_data, "y_negative"): - hard_negative_idx: torch.Tensor = torch.cat( - list(main_data.y_negative.values()) - ).to(device) + hard_negative_idx: torch.Tensor = main_data.y_negative[1] else: - hard_negative_idx = torch.empty(0, dtype=torch.long).to(device) + hard_negative_idx = torch.empty(0, dtype=torch.long, device=device) # Use local IDs to get the corresponding embeddings in the tensors diff --git a/gigl/distributed/dist_ablp_neighborloader.py b/gigl/distributed/dist_ablp_neighborloader.py index 50f42f5a9..c7224cb47 100644 --- a/gigl/distributed/dist_ablp_neighborloader.py +++ b/gigl/distributed/dist_ablp_neighborloader.py @@ -1,6 +1,7 @@ -from collections import abc, defaultdict +import warnings +from collections import abc from itertools import count -from typing import Optional, Union +from typing import Literal, Optional, Union import torch from graphlearn_torch.channel import SampleMessage @@ -27,6 +28,10 @@ SamplerOptions, resolve_sampler_options, ) +from gigl.distributed.utils.ablp import ( + label_edge_index_to_dict, + remap_labels_to_local_edge_indices, +) from gigl.distributed.utils.neighborloader import ( DatasetSchema, SamplingClusterSetup, @@ -43,9 +48,6 @@ from gigl.types.graph import ( DEFAULT_HOMOGENEOUS_EDGE_TYPE, DEFAULT_HOMOGENEOUS_NODE_TYPE, - label_edge_type_to_message_passing_edge_type, - message_passing_to_negative_label, - message_passing_to_positive_label, reverse_edge_type, select_label_edge_types, ) @@ -55,6 +57,190 @@ logger = Logger() +def _is_integral_tensor(tensor: torch.Tensor) -> bool: + return tensor.dtype in ( + torch.uint8, + torch.int8, + torch.int16, + torch.int32, + torch.int64, + ) + + +def _convert_graph_store_ablp_inputs( + input_nodes: dict[int, ABLPInputNodes], + num_storage_nodes: int, + edge_types: Optional[list[EdgeType]], + edge_dir: Literal["in", "out"], +) -> tuple[ + list[ABLPNodeSamplerInput], + NodeType, + list[EdgeType], + list[EdgeType], + list[EdgeType], +]: + """Validate GraphStore ABLP inputs and convert them to sampler inputs.""" + if not input_nodes: + raise ValueError("Graph Store ABLP input_nodes must be non-empty.") + + server_ranks = input_nodes.keys() + if min(server_ranks) < 0 or max(server_ranks) >= num_storage_nodes: + raise ValueError( + "When using Graph Store mode, server ranks must be in range " + f"[0, {num_storage_nodes}), received {list(server_ranks)}." + ) + if edge_types is None: + raise ValueError("Graph Store ABLP requires registered label edge types.") + + first_input = next(iter(input_nodes.values())) + input_type = first_input.anchor_node_type + supervision_edge_types = list(first_input.labels.keys()) + if not supervision_edge_types: + raise ValueError("Graph Store ABLP labels must be non-empty.") + + anchor_endpoint_index = 0 if edge_dir == "out" else 2 + invalid_supervision_edge_types = [ + edge_type + for edge_type in supervision_edge_types + if edge_type[anchor_endpoint_index] != input_type + ] + if invalid_supervision_edge_types: + raise ValueError( + f"For edge_dir={edge_dir!r}, registered supervision edge types must " + f"have anchor node type {input_type!r} at index " + f"{anchor_endpoint_index}; received {invalid_supervision_edge_types}." + ) + + positive_label_edge_types: list[EdgeType] = [] + negative_label_edge_types: list[EdgeType] = [] + negative_label_edge_type_by_supervision_edge_type: dict[ + EdgeType, Optional[EdgeType] + ] = {} + for supervision_edge_type in supervision_edge_types: + ( + positive_label_edge_type, + negative_label_edge_type, + ) = select_label_edge_types(supervision_edge_type, edge_types) + positive_label_edge_types.append(positive_label_edge_type) + negative_label_edge_type_by_supervision_edge_type[supervision_edge_type] = ( + negative_label_edge_type + ) + if negative_label_edge_type is not None: + negative_label_edge_types.append(negative_label_edge_type) + + expected_supervision_edge_type_set = set(supervision_edge_types) + for server_rank, ablp_input in input_nodes.items(): + if ablp_input.anchor_node_type != input_type: + raise ValueError( + f"All Graph Store ABLP inputs must have anchor node type " + f"{input_type!r}; server {server_rank} has " + f"{ablp_input.anchor_node_type!r}." + ) + if set(ablp_input.labels.keys()) != expected_supervision_edge_type_set: + raise ValueError( + "All Graph Store ABLP inputs must have the same supervision " + f"edge type keys {supervision_edge_types}; server {server_rank} " + f"has {list(ablp_input.labels.keys())}." + ) + + anchors = ablp_input.anchor_nodes + if anchors.dim() != 1 or not _is_integral_tensor(anchors): + raise ValueError( + f"Server {server_rank} anchor nodes must be a 1-D integral " + f"tensor, got shape {tuple(anchors.shape)} and dtype " + f"{anchors.dtype}." + ) + for supervision_edge_type in supervision_edge_types: + positive_labels, negative_labels = ablp_input.labels[supervision_edge_type] + if positive_labels is None: + raise ValueError( + f"Server {server_rank} positive labels for " + f"{supervision_edge_type} must be a tensor." + ) + for label_kind, label_tensor in ( + ("positive", positive_labels), + ("negative", negative_labels), + ): + if label_tensor is None: + continue + if label_tensor.dim() != 2 or not _is_integral_tensor(label_tensor): + raise ValueError( + f"Server {server_rank} {label_kind} labels for " + f"{supervision_edge_type} must be a 2-D integral tensor, " + f"got shape {tuple(label_tensor.shape)} and dtype " + f"{label_tensor.dtype}." + ) + if label_tensor.size(0) != anchors.numel(): + raise ValueError( + f"Server {server_rank} {label_kind} labels for " + f"{supervision_edge_type} have {label_tensor.size(0)} " + f"rows for {anchors.numel()} anchors." + ) + + expects_negative_labels = ( + negative_label_edge_type_by_supervision_edge_type[supervision_edge_type] + is not None + ) + if expects_negative_labels != (negative_labels is not None): + expected_description = "a tensor" if expects_negative_labels else "None" + raise ValueError( + f"Server {server_rank} negative labels for " + f"{supervision_edge_type} must be {expected_description} " + "to match registered topology." + ) + + input_data: list[ABLPNodeSamplerInput] = [] + for server_rank in range(num_storage_nodes): + if server_rank in input_nodes: + ablp_input = input_nodes[server_rank] + anchors = ablp_input.anchor_nodes + positive_label_by_edge_type = { + positive_label_edge_type: ablp_input.labels[supervision_edge_type][0] + for supervision_edge_type, positive_label_edge_type in zip( + supervision_edge_types, positive_label_edge_types + ) + } + negative_label_by_edge_type: dict[EdgeType, torch.Tensor] = {} + for supervision_edge_type in supervision_edge_types: + negative_label_edge_type = ( + negative_label_edge_type_by_supervision_edge_type[ + supervision_edge_type + ] + ) + if negative_label_edge_type is None: + continue + negative_labels = ablp_input.labels[supervision_edge_type][1] + assert negative_labels is not None + negative_label_by_edge_type[negative_label_edge_type] = negative_labels + else: + anchors = torch.empty(0, dtype=torch.long) + positive_label_by_edge_type = { + edge_type: torch.empty((0, 0), dtype=torch.long) + for edge_type in positive_label_edge_types + } + negative_label_by_edge_type = { + edge_type: torch.empty((0, 0), dtype=torch.long) + for edge_type in negative_label_edge_types + } + + input_data.append( + ABLPNodeSamplerInput( + node=anchors, + input_type=input_type, + positive_label_by_edge_types=positive_label_by_edge_type, + negative_label_by_edge_types=negative_label_by_edge_type, + ) + ) + + return ( + input_data, + input_type, + supervision_edge_types, + positive_label_edge_types, + negative_label_edge_types, + ) + + class DistABLPLoader(BaseDistLoader): # Counts instantiations of this class, per process. # This is needed so we can generate unique worker key for each instance, for graph store mode. @@ -90,17 +276,19 @@ def __init__( local_process_rank: Optional[int] = None, # TODO: (svij) Deprecate this local_process_world_size: Optional[int] = None, # TODO: (svij) Deprecate this non_blocking_transfers: bool = True, + use_edge_index_output: bool = False, ): """ Neighbor loader for Anchor Based Link Prediction (ABLP) tasks. - Note that for this class, the dataset must *always* be heterogeneous, - as we need separate edge types for positive and negative labels. + The dataset must *always* be heterogeneous here, since positive and + negative labels are carried as separate edge types. By default, the loader will return {py:class} `torch_geometric.data.HeteroData` (heterogeneous) objects, but will return a {py:class}`torch_geometric.data.Data` (homogeneous) object if the dataset is "labeled homogeneous". - The following fields may also be present: + The following fields may also be present (this describes the deprecated + default shape; see ``use_edge_index_output`` for the tensor format): - `y_positive`: `dict[int, torch.Tensor]` mapping from local anchor node id to a tensor of positive label node ids. - `y_negative`: (Optional) `dict[int, torch.Tensor]` mapping from local anchor node id to a tensor of negative @@ -138,6 +326,16 @@ def __init__( - `y_positive`: {(a, to, b): {0: torch.tensor([1])}, (a, to, c): {0: torch.tensor([2])}} - `y_negative`: {(a, to, b): {0: torch.tensor([3])}, (a, to, c): {0: torch.tensor([4])}} + With ``use_edge_index_output=True``, each label field is a ``[2, E]`` + tensor. Row 0 contains local anchor indices and row 1 contains local + label-node indices. For the example above: + + - ``y_positive = torch.tensor([[0], [1]])`` + - ``y_negative = torch.tensor([[0], [2]])`` + + Multiple supervision edge types produce ``dict[EdgeType, torch.Tensor]``. + Pair order within an anchor is unspecified. + Args: dataset (Union[DistDataset, RemoteDistDataset]): The dataset to sample from. If this is a `RemoteDistDataset`, then we are in "Graph Store" mode. @@ -213,11 +411,23 @@ def __init__( is used instead. See https://docs.pytorch.org/tutorials/intermediate/pinmem_nonblock.html for background on pinned memory and non-blocking transfers. + use_edge_index_output (bool): Return labels as ``[2, E]`` edge-index + tensors instead of the deprecated ragged dictionaries. Row 0 + contains local anchor indices and row 1 contains local label-node + indices. Defaults to ``False`` for backward compatibility. """ # Set self._shutdowned right away, that way if we throw here, and __del__ is called, # then we can properly clean up and don't get extraneous error messages. self._shutdowned = True + if not use_edge_index_output: + warnings.warn( + "The ragged dictionary ABLP label output is deprecated. Pass " + "use_edge_index_output=True to receive [2, E] label edge indices.", + FutureWarning, + stacklevel=2, + ) + self._use_edge_index_output = use_edge_index_output sampler_options = resolve_sampler_options(num_neighbors, sampler_options) @@ -612,9 +822,57 @@ def _setup_for_graph_store( Tuple of (list[ABLPNodeSamplerInput], RemoteDistSamplingWorkerOptions, DatasetSchema, backend_key). """ + edge_types = dataset.fetch_edge_types() + edge_dir = dataset.fetch_edge_dir() + if edge_dir == "in": + evaluated_edge_dir: Literal["in", "out"] = "in" + elif edge_dir == "out": + evaluated_edge_dir = "out" + else: + raise ValueError( + f"Expected GraphStore edge_dir to be 'in' or 'out', got {edge_dir!r}." + ) + + validation_error: Optional[str] = None + converted_inputs: Optional[ + tuple[ + list[ABLPNodeSamplerInput], + NodeType, + list[EdgeType], + list[EdgeType], + list[EdgeType], + ] + ] = None + try: + converted_inputs = _convert_graph_store_ablp_inputs( + input_nodes=input_nodes, + num_storage_nodes=dataset.cluster_info.num_storage_nodes, + edge_types=edge_types, + edge_dir=evaluated_edge_dir, + ) + except ValueError as error: + validation_error = str(error) + + validation_errors: list[Optional[str]] = [ + None + ] * torch.distributed.get_world_size() + torch.distributed.all_gather_object(validation_errors, validation_error) + if any(error is not None for error in validation_errors): + raise ValueError( + "Graph Store ABLP input validation failed across compute ranks: " + f"{validation_errors}" + ) + assert converted_inputs is not None + ( + input_data, + input_type, + self._supervision_edge_types, + self._positive_label_edge_types, + self._negative_label_edge_types, + ) = converted_inputs + node_feature_info = dataset.fetch_node_feature_info() edge_feature_info = dataset.fetch_edge_feature_info() - edge_types = dataset.fetch_edge_types() compute_rank = torch.distributed.get_rank() backend_key = f"dist_ablp_loader_{self._instance_count}" worker_key = f"{backend_key}_compute_rank_{compute_rank}" @@ -632,107 +890,13 @@ def _setup_for_graph_store( f"tcp://{worker_options.master_addr}:{worker_options.master_port}" ) - # Validate server ranks - servers = input_nodes.keys() - if len(servers) > 0: - if ( - max(servers) >= dataset.cluster_info.num_storage_nodes - or min(servers) < 0 - ): - raise ValueError( - f"When using Graph Store mode, the server ranks must be in range " - f"[0, {dataset.cluster_info.num_storage_nodes}), " - f"received inputs for servers: {list(servers)}" - ) - - # Extract node type and label edge types from the ABLPInputNodes dataclass. - # All entries should have the same anchor_node_type and edge type keys. - first_input = next(iter(input_nodes.values())) - input_type = first_input.anchor_node_type is_homogeneous_with_labeled_edge_type = ( input_type == DEFAULT_HOMOGENEOUS_NODE_TYPE ) - # Extract supervision edge types and derive label edge types from the - # ABLPInputNodes.labels dict (keyed by supervision edge type). - self._supervision_edge_types = list(first_input.labels.keys()) - has_negatives = False - for ablp_input in input_nodes.values(): - for maybe_negative_labels in ablp_input.labels.values(): - if maybe_negative_labels is not None: - has_negatives = True - break - - self._positive_label_edge_types = [ - message_passing_to_positive_label(et) for et in self._supervision_edge_types - ] - self._negative_label_edge_types = ( - [ - message_passing_to_negative_label(et) - for et in self._supervision_edge_types - ] - if has_negatives - else [] - ) - - # Graph Store mode currently only supports a single supervision edge type, - # so the labels dict must have exactly 1 entry. - if len(self._supervision_edge_types) != 1: - raise ValueError( - f"Graph Store mode currently only supports a single supervision edge type, " - f"but ABLPInputNodes.labels has {len(self._supervision_edge_types)} " - f"entries: {self._supervision_edge_types}" - ) - logger.info(f"Positive label edge types: {self._positive_label_edge_types}") logger.info(f"Negative label edge types: {self._negative_label_edge_types}") - # Convert from ABLPInputNodes to list of ABLPNodeSamplerInput (one per server) - input_data: list[ABLPNodeSamplerInput] = [] - for server_rank in range(dataset.cluster_info.num_storage_nodes): - positive_label_by_edge_type: dict[EdgeType, torch.Tensor] = {} - negative_label_by_edge_type: dict[EdgeType, torch.Tensor] = {} - if server_rank in input_nodes: - ablp_input_nodes = input_nodes[server_rank] - anchors = ablp_input_nodes.anchor_nodes - for supervision_edge_type, ( - positive_labels, - negative_labels, - ) in ablp_input_nodes.labels.items(): - positive_label_by_edge_type[ - message_passing_to_positive_label(supervision_edge_type) - ] = positive_labels - if negative_labels is not None: - negative_label_by_edge_type[ - message_passing_to_negative_label(supervision_edge_type) - ] = negative_labels - else: - # Empty input for servers with no data for this rank - anchors = torch.empty(0, dtype=torch.long) - positive_label_by_edge_type = { - et: torch.empty(0, 0, dtype=torch.long) - for et in self._positive_label_edge_types - } - if has_negatives: - negative_label_by_edge_type = { - et: torch.empty(0, 0, dtype=torch.long) - for et in self._negative_label_edge_types - } - - logger.info( - f"Rank: {torch.distributed.get_rank()}! Building ABLPNodeSamplerInput for server rank: {server_rank} " - f"with input type: {input_type}. anchors: {anchors.shape}, " - f"positive_labels edge types: {list(positive_label_by_edge_type.keys())}, " - f"negative_labels edge types: {list(negative_label_by_edge_type.keys())}" - ) - ablp_input = ABLPNodeSamplerInput( - node=anchors, - input_type=input_type, - positive_label_by_edge_types=positive_label_by_edge_type, - negative_label_by_edge_types=negative_label_by_edge_type, - ) - input_data.append(ablp_input) - return ( input_data, worker_options, @@ -741,7 +905,7 @@ def _setup_for_graph_store( edge_types=edge_types, node_feature_info=node_feature_info, edge_feature_info=edge_feature_info, - edge_dir=dataset.fetch_edge_dir(), + edge_dir=evaluated_edge_dir, ), backend_key, ) @@ -752,82 +916,124 @@ def _set_labels( positive_labels_by_label_edge_type: dict[EdgeType, torch.Tensor], negative_labels_by_label_edge_type: dict[EdgeType, torch.Tensor], ) -> Union[Data, HeteroData]: - """ - Sets the labels and relevant fields in the torch_geometric Data object, converting the global node ids for labels to their - local index. Removes inserted supervision edge type from the data variables, since this is an implementation detail and should not be - exposed in the final HeteroData/Data object. + """Attach ABLP labels to the collated graph, remapped to subgraph-local indices. + + The tensor output uses ``[2, E]`` label edge indices. The compatibility + path converts the same tensors to the deprecated ragged dictionaries. + Singleton supervision retains the historical flattened output; with + multiple supervision types, positive and negative outputs remain keyed + by edge type even when only one type has negative labels. + Args: - data (Union[Data, HeteroData]): Graph to provide labels for - positive_labels_by_label_edge_type (dict[EdgeType, torch.Tensor]): Dict[positive label edge type, label ID tensor], - where the ith row of the tensor corresponds to the ith anchor node ID. - negative_labels_by_label_edge_type (dict[EdgeType, torch.Tensor]): Dict[negative label edge type, label ID tensor], - where the ith row of the tensor corresponds to the ith anchor node ID. + data (Union[Data, HeteroData]): Graph to attach labels to. + positive_labels_by_label_edge_type (dict[EdgeType, torch.Tensor]): Per + positive-label edge type, a ``[N_anchors, M]`` tensor whose ``i``-th + row holds the global label ids of the ``i``-th anchor. + negative_labels_by_label_edge_type (dict[EdgeType, torch.Tensor]): As + above, for negative-label edge types. + Returns: - Union[Data, HeteroData]: torch_geometric HeteroData/Data object with the filtered edge fields and labels set as properties of the instance + Union[Data, HeteroData]: The same object with the supervision edge fields + stripped and ``y_positive`` (and ``y_negative`` when present) attached. + + Raises: + ValueError: If no positive labels are found in ``data``. """ - # shape [N], where N is the number of nodes in the subgraph, and local_node_to_global_node[i] gives the global node id for local node id `i` - node_type_to_local_node_to_global_node: dict[NodeType, torch.Tensor] = {} + local_id_to_global_id_by_node_type: dict[NodeType, torch.Tensor] = {} if isinstance(data, HeteroData): - for e_type in self._supervision_edge_types: - node_type_to_local_node_to_global_node[e_type[0]] = data[e_type[0]].node - node_type_to_local_node_to_global_node[e_type[2]] = data[e_type[2]].node + # Key the map off the label edge types rather than + # self._supervision_edge_types: the latter is stored reversed when + # edge_dir="in" (see __init__), so its dst node type is the anchor + # type there, while label edge types always arrive in outward form. + # Deriving from the keys that _remap_labels_by_edge_type will look up + # keeps the two in step regardless of edge direction. + for label_edge_type, label_tensor in ( + *positive_labels_by_label_edge_type.items(), + *negative_labels_by_label_edge_type.items(), + ): + # Mirror the zero-anchor skip in _remap_labels_by_edge_type. An + # absent node type would otherwise be auto-created as an empty + # store by HeteroData.__getitem__, turning a clear KeyError into + # an AttributeError on .node. + if label_tensor.size(0) == 0: + continue + supervision_node_type = label_edge_type[2] + local_id_to_global_id_by_node_type[supervision_node_type] = data[ + supervision_node_type + ].node else: - node_type_to_local_node_to_global_node[DEFAULT_HOMOGENEOUS_NODE_TYPE] = ( + local_id_to_global_id_by_node_type[DEFAULT_HOMOGENEOUS_NODE_TYPE] = ( data.node ) - output_positive_labels: dict[EdgeType, dict[int, torch.Tensor]] = defaultdict( - dict - ) - output_negative_labels: dict[EdgeType, dict[int, torch.Tensor]] = defaultdict( - dict + + ( + positive_label_edge_indices, + negative_label_edge_indices, + ) = remap_labels_to_local_edge_indices( + local_id_to_global_id_by_node_type=local_id_to_global_id_by_node_type, + positive_labels_by_edge_type=positive_labels_by_label_edge_type, + negative_labels_by_edge_type=negative_labels_by_label_edge_type, ) - # We always have supervision edge types of the form (anchor_node_type, to, supervision_node_type) - # So we can index into the edge type accordingly. - edge_index = 2 - for edge_type, label_tensor in positive_labels_by_label_edge_type.items(): - for local_anchor_node_id in range(label_tensor.size(0)): - positive_mask = ( - node_type_to_local_node_to_global_node[ - edge_type[edge_index] - ].unsqueeze(1) - == label_tensor[local_anchor_node_id] - ) # shape [N, P], where N is the number of nodes and P is the number of positive labels for the current anchor node - - # Gets the indexes of the items in local_node_to_global_node which match any of the positive labels for the current anchor node - output_positive_labels[ - label_edge_type_to_message_passing_edge_type(edge_type) - ][local_anchor_node_id] = torch.nonzero(positive_mask)[:, 0].to( - self.to_device + output_positive_labels: dict[ + EdgeType, Union[torch.Tensor, dict[int, torch.Tensor]] + ] = {} + output_negative_labels: dict[ + EdgeType, Union[torch.Tensor, dict[int, torch.Tensor]] + ] = {} + if self._use_edge_index_output: + for edge_type, label_edge_index in positive_label_edge_indices.items(): + output_positive_labels[edge_type] = label_edge_index + for edge_type, label_edge_index in negative_label_edge_indices.items(): + output_negative_labels[edge_type] = label_edge_index + else: + label_edge_types = ( + positive_label_edge_indices.keys() | negative_label_edge_indices.keys() + ) + if isinstance(data, HeteroData): + num_anchors_by_edge_type = { + edge_type: int(data[edge_type[0]].batch_size) + for edge_type in label_edge_types + } + else: + num_anchors_by_edge_type = { + edge_type: int(data.batch_size) for edge_type in label_edge_types + } + output_positive_labels = { + edge_type: label_edge_index_to_dict( + label_edge_index=label_edge_index, + num_anchors=num_anchors_by_edge_type[edge_type], ) - # Shape [X], where X is the number of indexes in the original local_node_to_global_node which match a node in the positive labels for the current anchor node - - for edge_type, label_tensor in negative_labels_by_label_edge_type.items(): - for local_anchor_node_id in range(label_tensor.size(0)): - negative_mask = ( - node_type_to_local_node_to_global_node[ - edge_type[edge_index] - ].unsqueeze(1) - == label_tensor[local_anchor_node_id] - ) # shape [N, M], where N is the number of nodes and M is the number of negative labels for the current anchor node - - # Gets the indexes of the items in local_node_to_global_node which match any of the negative labels for the current anchor node - output_negative_labels[ - label_edge_type_to_message_passing_edge_type(edge_type) - ][local_anchor_node_id] = torch.nonzero(negative_mask)[:, 0].to( - self.to_device + for edge_type, label_edge_index in positive_label_edge_indices.items() + } + output_negative_labels = { + edge_type: label_edge_index_to_dict( + label_edge_index=label_edge_index, + num_anchors=num_anchors_by_edge_type[edge_type], ) - # Shape [X], where X is the number of indexes in the original local_node_to_global_node which match a node in the negative labels for the current anchor node + for edge_type, label_edge_index in negative_label_edge_indices.items() + } if not output_positive_labels: raise ValueError("No positive labels were found in the data!") - elif len(output_positive_labels) == 1: + elif len(self._supervision_edge_types) == 1: + if len(output_positive_labels) != 1: + raise ValueError( + "Expected exactly one positive output for a single " + f"supervision edge type, got {list(output_positive_labels)}." + ) data.y_positive = next(iter(output_positive_labels.values())) else: data.y_positive = output_positive_labels - if len(output_negative_labels) == 1: - data.y_negative = next(iter(output_negative_labels.values())) - elif len(output_negative_labels) > 0: - data.y_negative = output_negative_labels + if output_negative_labels: + if len(self._supervision_edge_types) == 1: + if len(output_negative_labels) != 1: + raise ValueError( + "Expected at most one negative output for a single " + f"supervision edge type, got {list(output_negative_labels)}." + ) + data.y_negative = next(iter(output_negative_labels.values())) + else: + data.y_negative = output_negative_labels return data def _collate_fn(self, msg: SampleMessage) -> Union[Data, HeteroData]: diff --git a/gigl/distributed/graph_store/dist_server.py b/gigl/distributed/graph_store/dist_server.py index 3c8a87ee1..b19057fe8 100644 --- a/gigl/distributed/graph_store/dist_server.py +++ b/gigl/distributed/graph_store/dist_server.py @@ -96,6 +96,7 @@ def compute_process(): from gigl.src.common.types.graph_data import EdgeType, NodeType from gigl.types.graph import FeatureInfo, select_label_edge_types from gigl.utils.data_splitters import get_labels_for_anchor_nodes +from gigl.utils.sampling import ABLPInputNodes from gigl.utils.share_memory import share_memory SERVER_EXIT_STATUS_CHECK_INTERVAL = 5.0 @@ -542,38 +543,69 @@ def get_node_types(self) -> Optional[list[NodeType]]: def get_ablp_input( self, request: FetchABLPInputRequest, - ) -> tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + ) -> ABLPInputNodes: """Get the ABLP (Anchor Based Link Prediction) input for distributed processing. Args: request: The ABLP fetch request, including split, node type, - supervision edge type, and an optional contiguous server slice. + registered supervision edge types, and an optional contiguous + server slice. Returns: - A tuple containing the anchor nodes for the rank, the positive labels, and the negative labels. - The positive labels are of shape [N, M], where N is the number of anchor nodes and M is the number of positive labels. - The negative labels are of shape [N, M], where N is the number of anchor nodes and M is the number of negative labels. - The negative labels may be None if no negative labels are available. + An ABLPInputNodes containing the anchors and one positive/optional + negative label pair for every requested supervision edge type. Raises: - ValueError: If the split is invalid. + ValueError: If the split or supervision schema is invalid. """ + if not request.supervision_edge_types: + raise ValueError("supervision_edge_types must be non-empty.") + if len(set(request.supervision_edge_types)) != len( + request.supervision_edge_types + ): + raise ValueError( + "supervision_edge_types must not contain duplicates, received " + f"{request.supervision_edge_types}." + ) + + anchor_endpoint_index = 0 if self.dataset.edge_dir == "out" else 2 + invalid_supervision_edge_types = [ + edge_type + for edge_type in request.supervision_edge_types + if edge_type[anchor_endpoint_index] != request.node_type + ] + if invalid_supervision_edge_types: + raise ValueError( + f"For edge_dir={self.dataset.edge_dir!r}, registered supervision " + f"edge types must have anchor node type {request.node_type!r} at " + f"index {anchor_endpoint_index}; received invalid edge types " + f"{invalid_supervision_edge_types}." + ) + anchors = self._get_node_ids( split=request.split, node_type=request.node_type, server_slice=request.server_slice, ) - positive_label_edge_type, negative_label_edge_type = select_label_edge_types( - request.supervision_edge_type, self.dataset.get_edge_types() - ) - positive_labels, negative_labels = get_labels_for_anchor_nodes( - self.dataset, - anchors, - positive_label_edge_type, - negative_label_edge_type, - max_labels_per_anchor_node=self.dataset.max_labels_per_anchor_node, + labels: dict[EdgeType, tuple[torch.Tensor, Optional[torch.Tensor]]] = {} + edge_types = self.dataset.get_edge_types() + for supervision_edge_type in request.supervision_edge_types: + ( + positive_label_edge_type, + negative_label_edge_type, + ) = select_label_edge_types(supervision_edge_type, edge_types) + labels[supervision_edge_type] = get_labels_for_anchor_nodes( + self.dataset, + anchors, + positive_label_edge_type, + negative_label_edge_type, + max_labels_per_anchor_node=self.dataset.max_labels_per_anchor_node, + ) + return ABLPInputNodes( + anchor_nodes=anchors, + anchor_node_type=request.node_type, + labels=labels, ) - return anchors, positive_labels, negative_labels def init_sampling_backend(self, opts: InitSamplingBackendRequest) -> int: """Create or reuse a shared sampling backend for one loader instance. diff --git a/gigl/distributed/graph_store/messages.py b/gigl/distributed/graph_store/messages.py index 6c659cc9a..87c28039d 100644 --- a/gigl/distributed/graph_store/messages.py +++ b/gigl/distributed/graph_store/messages.py @@ -94,7 +94,8 @@ class FetchABLPInputRequest: Args: split: The split of the dataset to get ABLP input from. node_type: The type of anchor nodes to retrieve. - supervision_edge_type: The edge type used for supervision. + supervision_edge_types: The registered sampling-orientation edge types + used for supervision. server_slice: An optional :class:`~gigl.distributed.graph_store.sharding.ServerSlice` describing the fraction of this server's data to return. When ``None``, all of the server's data is returned. @@ -102,16 +103,20 @@ class FetchABLPInputRequest: Examples: Fetch training ABLP input without sharding: - >>> FetchABLPInputRequest(split="train", node_type="user", supervision_edge_type=("user", "to", "item")) + >>> FetchABLPInputRequest( + ... split="train", + ... node_type="user", + ... supervision_edge_types=(("user", "to", "item"),), + ... ) Fetch with a server slice: >>> FetchABLPInputRequest(split="train", node_type="user", - ... supervision_edge_type=("user", "to", "item"), + ... supervision_edge_types=(("user", "to", "item"),), ... server_slice=ServerSlice(0, 0, 1, 2)) """ split: Union[Literal["train", "val", "test"], str] node_type: NodeType - supervision_edge_type: EdgeType + supervision_edge_types: tuple[EdgeType, ...] server_slice: Optional[ServerSlice] = None diff --git a/gigl/distributed/graph_store/remote_dist_dataset.py b/gigl/distributed/graph_store/remote_dist_dataset.py index e58784dc0..8afa2383d 100644 --- a/gigl/distributed/graph_store/remote_dist_dataset.py +++ b/gigl/distributed/graph_store/remote_dist_dataset.py @@ -21,12 +21,89 @@ DEFAULT_HOMOGENEOUS_EDGE_TYPE, DEFAULT_HOMOGENEOUS_NODE_TYPE, FeatureInfo, + reverse_edge_type, + select_label_edge_types, ) from gigl.utils.sampling import ABLPInputNodes logger = Logger() +def _resolve_registered_supervision_edge_types( + supervision_edge_types: list[EdgeType], + anchor_node_type: NodeType, + edge_dir: Literal["in", "out"], + registered_edge_types: list[EdgeType], +) -> tuple[tuple[EdgeType, ...], set[EdgeType]]: + """Resolve public ABLP supervision types to registered sampling orientation. + + Public supervision edge types use the colocated anchor-outward convention. + For incoming graphs, the legacy GraphStore convention (anchor as destination) + remains accepted when every supplied edge type uses it. + + Args: + supervision_edge_types: Non-empty public supervision edge types. + anchor_node_type: Node type of the anchor nodes. + edge_dir: Registered graph sampling direction. + registered_edge_types: Edge types registered in GraphStore. + + Returns: + A tuple containing the ordered registered supervision edge types and the + subset with registered negative label topology. + + Raises: + ValueError: If the edge types are duplicated, use inconsistent + orientations, do not share the anchor node type, or lack positive + label topology. + """ + if not supervision_edge_types: + raise ValueError("supervision_edge_type must be a non-empty list.") + if len(set(supervision_edge_types)) != len(supervision_edge_types): + raise ValueError( + "supervision_edge_type must not contain duplicates, received " + f"{supervision_edge_types}." + ) + + all_anchor_outward = all( + edge_type[0] == anchor_node_type for edge_type in supervision_edge_types + ) + if edge_dir == "out": + if not all_anchor_outward: + raise ValueError( + "For edge_dir='out', all supervision edge types must be " + f"anchor-outward from {anchor_node_type!r}, received " + f"{supervision_edge_types}." + ) + resolved_edge_types = tuple(supervision_edge_types) + else: + all_registered_incoming = all( + edge_type[2] == anchor_node_type for edge_type in supervision_edge_types + ) + if all_anchor_outward: + resolved_edge_types = tuple( + reverse_edge_type(edge_type) for edge_type in supervision_edge_types + ) + elif all_registered_incoming: + # Backward compatibility for the existing GraphStore convention. + resolved_edge_types = tuple(supervision_edge_types) + else: + raise ValueError( + "For edge_dir='in', supervision edge types must be uniformly " + f"anchor-outward from {anchor_node_type!r} or uniformly use the " + "legacy registered incoming orientation with that anchor as " + f"destination; received {supervision_edge_types}." + ) + + edge_types_with_negatives: set[EdgeType] = set() + for resolved_edge_type in resolved_edge_types: + _, negative_label_edge_type = select_label_edge_types( + resolved_edge_type, registered_edge_types + ) + if negative_label_edge_type is not None: + edge_types_with_negatives.add(resolved_edge_type) + return resolved_edge_types, edge_types_with_negatives + + class RemoteDistDataset: def __init__( self, @@ -381,10 +458,16 @@ def _fetch_ablp_input( self, split: Literal["train", "val", "test"], node_type: NodeType = DEFAULT_HOMOGENEOUS_NODE_TYPE, - supervision_edge_type: EdgeType = DEFAULT_HOMOGENEOUS_EDGE_TYPE, + supervision_edge_types: tuple[EdgeType, ...] = (DEFAULT_HOMOGENEOUS_EDGE_TYPE,), + supervision_edge_types_with_negatives: Optional[set[EdgeType]] = None, assignments: Optional[dict[int, ServerSlice]] = None, - ) -> dict[int, tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]]: + ) -> dict[int, ABLPInputNodes]: """Fetches ABLP input from the storage nodes for the current compute node (machine).""" + evaluated_edge_types_with_negatives = ( + supervision_edge_types_with_negatives + if supervision_edge_types_with_negatives is not None + else set() + ) # Build per-server requests requests: dict[int, FetchABLPInputRequest] = {} if assignments is None: @@ -393,14 +476,14 @@ def _fetch_ablp_input( requests[server_rank] = FetchABLPInputRequest( split=split, node_type=node_type, - supervision_edge_type=supervision_edge_type, + supervision_edge_types=supervision_edge_types, ) else: for server_rank, server_slice in assignments.items(): requests[server_rank] = FetchABLPInputRequest( split=split, node_type=node_type, - supervision_edge_type=supervision_edge_type, + supervision_edge_types=supervision_edge_types, server_slice=server_slice, ) @@ -408,31 +491,32 @@ def _fetch_ablp_input( logger.info( f"Fetching ABLP input (sharded={sharded}) " f"with node type {node_type}, split {split}, and " - f"supervision edge type {supervision_edge_type}. " + f"registered supervision edge types {supervision_edge_types}. " f"Requesting from servers: {sorted(requests.keys())}" ) # Dispatch all futures - futures: dict[ - int, - torch.futures.Future[ - tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]] - ], - ] = { + futures: dict[int, torch.futures.Future[ABLPInputNodes]] = { server_rank: async_request_server( server_rank, DistServer.get_ablp_input, request ) for server_rank, request in requests.items() } - def _empty_ablp_result() -> tuple[ - torch.Tensor, torch.Tensor, Optional[torch.Tensor] - ]: - """Return an empty ABLP result tuple: (anchor_nodes, positive_labels, negative_labels).""" - return ( - torch.empty(0, dtype=torch.long), - torch.empty((0, 0), dtype=torch.long), - None, + def _empty_ablp_result() -> ABLPInputNodes: + """Return schema-complete empty ABLP input for an unassigned server.""" + return ABLPInputNodes( + anchor_nodes=torch.empty(0, dtype=torch.long), + anchor_node_type=node_type, + labels={ + edge_type: ( + torch.empty((0, 0), dtype=torch.long), + torch.empty((0, 0), dtype=torch.long) + if edge_type in evaluated_edge_types_with_negatives + else None, + ) + for edge_type in supervision_edge_types + }, ) # Collect results, filling empty tuples for unrequested servers @@ -450,7 +534,7 @@ def fetch_ablp_input( rank: Optional[int] = None, world_size: Optional[int] = None, anchor_node_type: Optional[NodeType] = None, - supervision_edge_type: Optional[EdgeType] = None, + supervision_edge_type: Optional[Union[EdgeType, list[EdgeType]]] = None, ) -> dict[int, ABLPInputNodes]: """Fetch ABLP (Anchor Based Link Prediction) input from the storage nodes. @@ -474,7 +558,9 @@ def fetch_ablp_input( anchor_node_type: The type of the anchor nodes to retrieve. Must be provided for heterogeneous graphs. Must be ``None`` for labeled homogeneous graphs. - supervision_edge_type: The edge type for supervision. + supervision_edge_type: One or more anchor-outward edge types for + supervision. Incoming GraphStore's legacy registered orientation + remains accepted when all types use it. Must be provided for heterogeneous graphs. Must be ``None`` for labeled homogeneous graphs. @@ -484,8 +570,11 @@ def fetch_ablp_input( - ``anchor_node_type``: The node type of the anchor nodes, or ``DEFAULT_HOMOGENEOUS_NODE_TYPE`` for labeled homogeneous. - ``anchor_nodes``: 1D tensor of anchor node IDs for the split. - - ``positive_labels``: Dict mapping positive label EdgeType to a 2D tensor [N, M]. - - ``negative_labels``: Optional dict mapping negative label EdgeType to a 2D tensor [N, M]. + - ``labels``: An ordered dict mapping each registered-orientation + supervision edge type to a positive 2D tensor and an optional + negative 2D tensor, both shaped ``[N, M]``. For incoming graphs, + canonical public types are reversed here and reversed back during + loader collation. Raises: ValueError: If only one of ``rank`` or ``world_size`` is provided. @@ -552,38 +641,67 @@ def fetch_ablp_input( else: evaluated_anchor_node_type = anchor_node_type if supervision_edge_type is None: - evaluated_supervision_edge_type = DEFAULT_HOMOGENEOUS_EDGE_TYPE + evaluated_supervision_edge_types = [DEFAULT_HOMOGENEOUS_EDGE_TYPE] + elif isinstance(supervision_edge_type, list): + evaluated_supervision_edge_types = list(supervision_edge_type) else: - evaluated_supervision_edge_type = supervision_edge_type + evaluated_supervision_edge_types = [supervision_edge_type] del anchor_node_type, supervision_edge_type + if not evaluated_supervision_edge_types: + raise ValueError("supervision_edge_type must be a non-empty list.") + if len(set(evaluated_supervision_edge_types)) != len( + evaluated_supervision_edge_types + ): + raise ValueError( + "supervision_edge_type must not contain duplicates, received " + f"{evaluated_supervision_edge_types}." + ) + if ( + evaluated_anchor_node_type == DEFAULT_HOMOGENEOUS_NODE_TYPE + and evaluated_supervision_edge_types != [DEFAULT_HOMOGENEOUS_EDGE_TYPE] + ): + raise ValueError( + "Labeled homogeneous GraphStore input supports only the default " + "homogeneous supervision edge type." + ) + assignments = self._compute_assignments_if_needed( rank=rank, world_size=world_size, ) + edge_dir = self.fetch_edge_dir() + if edge_dir == "in": + evaluated_edge_dir: Literal["in", "out"] = "in" + elif edge_dir == "out": + evaluated_edge_dir = "out" + else: + raise ValueError( + f"Expected GraphStore edge_dir to be 'in' or 'out', got {edge_dir!r}." + ) + registered_edge_types = self.fetch_edge_types() + if registered_edge_types is None: + raise ValueError( + "ABLP GraphStore input requires registered label edge types." + ) + ( + evaluated_supervision_edge_types, + supervision_edge_types_with_negatives, + ) = _resolve_registered_supervision_edge_types( + supervision_edge_types=evaluated_supervision_edge_types, + anchor_node_type=evaluated_anchor_node_type, + edge_dir=evaluated_edge_dir, + registered_edge_types=registered_edge_types, + ) + raw_inputs = self._fetch_ablp_input( split=split, node_type=evaluated_anchor_node_type, - supervision_edge_type=evaluated_supervision_edge_type, + supervision_edge_types=evaluated_supervision_edge_types, + supervision_edge_types_with_negatives=supervision_edge_types_with_negatives, assignments=assignments, ) - return { - server_rank: ABLPInputNodes( - anchor_node_type=evaluated_anchor_node_type, - anchor_nodes=anchors, - labels={ - evaluated_supervision_edge_type: ( - positive_labels, - negative_labels, - ) - }, - ) - for server_rank, ( - anchors, - positive_labels, - negative_labels, - ) in raw_inputs.items() - } + return raw_inputs def fetch_edge_types(self) -> Optional[list[EdgeType]]: """Fetch the edge types from the registered dataset. diff --git a/gigl/distributed/utils/ablp.py b/gigl/distributed/utils/ablp.py new file mode 100644 index 000000000..aeff1e569 --- /dev/null +++ b/gigl/distributed/utils/ablp.py @@ -0,0 +1,160 @@ +"""Utilities for remapping ABLP labels to sampled-subgraph indices.""" + +import torch +from torch_geometric.typing import EdgeType + +from gigl.src.common.types.graph_data import NodeType +from gigl.types.graph import label_edge_type_to_message_passing_edge_type +from gigl.utils.data_splitters import PADDING_NODE + + +def label_edge_index_to_dict( + label_edge_index: torch.Tensor, num_anchors: int +) -> dict[int, torch.Tensor]: + """Convert a label edge index to the deprecated per-anchor dictionary. + + The edge index must be grouped by anchor, as guaranteed by + :func:`remap_labels_to_local_edge_indices`. + + Args: + label_edge_index: A ``[2, E]`` tensor. Row 0 contains local anchor + indices and row 1 contains local label-node indices. + num_anchors: Number of anchors in the sampled batch, including anchors + with no labels. + + Returns: + A mapping from every local anchor index to its local label-node indices. + """ + counts = torch.bincount(label_edge_index[0], minlength=num_anchors) + labels_by_anchor = torch.split(label_edge_index[1], counts.tolist()) + return {anchor: labels_by_anchor[anchor] for anchor in range(num_anchors)} + + +def _remap_label_tensor_to_local_edge_index( + label_tensor: torch.Tensor, + sorted_global_node_ids: torch.Tensor, + local_ids_by_sorted_global_id: torch.Tensor, +) -> torch.Tensor: + """Remap one padded global-label tensor to a local label edge index. + + Args: + label_tensor: A ``[N_anchors, M]`` tensor of global label-node ids, + padded with ``PADDING_NODE``. + sorted_global_node_ids: Sorted global ids for the sampled supervision + nodes. + local_ids_by_sorted_global_id: Local node indices in the order given by + ``sorted_global_node_ids``. + + Returns: + A ``[2, E]`` long tensor on the input device. Row 0 contains local + anchor indices and row 1 contains local label-node indices. Labels not + present in the sampled subgraph are omitted. + + Raises: + ValueError: If the sampled node map contains duplicate global ids. + """ + num_anchors = int(label_tensor.size(0)) + num_nodes = int(sorted_global_node_ids.size(0)) + empty_edge_index = torch.empty((2, 0), dtype=torch.long, device=label_tensor.device) + if num_anchors == 0: + return empty_edge_index + + num_labels = int(label_tensor.size(1)) + candidate_global_label_ids = label_tensor.reshape(-1) + candidate_anchor_indices = torch.arange( + num_anchors, device=label_tensor.device + ).repeat_interleave(num_labels) + + is_not_padding = candidate_global_label_ids != PADDING_NODE + candidate_global_label_ids = candidate_global_label_ids[is_not_padding] + candidate_anchor_indices = candidate_anchor_indices[is_not_padding] + if num_nodes == 0 or candidate_global_label_ids.numel() == 0: + return empty_edge_index + + if not bool((sorted_global_node_ids[1:] > sorted_global_node_ids[:-1]).all()): + raise ValueError( + "Vectorized label remapping requires unique global ids in the " + "sampled node map." + ) + + insertion_indices = torch.searchsorted( + sorted_global_node_ids, candidate_global_label_ids + ) + # An absent label larger than every sampled id inserts at num_nodes. Clamp + # before gathering; the exact-match mask below still discards that label. + insertion_indices = insertion_indices.clamp_(max=num_nodes - 1) + is_in_sampled_subgraph = ( + sorted_global_node_ids[insertion_indices] == candidate_global_label_ids + ) + + local_label_indices = local_ids_by_sorted_global_id[insertion_indices][ + is_in_sampled_subgraph + ] + anchor_indices = candidate_anchor_indices[is_in_sampled_subgraph] + return torch.stack((anchor_indices, local_label_indices)) + + +def _remap_labels_by_edge_type( + labels_by_edge_type: dict[EdgeType, torch.Tensor], + local_id_to_global_id_by_node_type: dict[NodeType, torch.Tensor], + sorted_node_lookup_by_type: dict[NodeType, tuple[torch.Tensor, torch.Tensor]], +) -> dict[EdgeType, torch.Tensor]: + """Remap positive or negative labels for each supervision edge type.""" + output: dict[EdgeType, torch.Tensor] = {} + for label_edge_type, label_tensor in labels_by_edge_type.items(): + if label_tensor.size(0) == 0: + continue + + supervision_node_type = label_edge_type[2] + if supervision_node_type not in sorted_node_lookup_by_type: + sorted_node_lookup_by_type[supervision_node_type] = torch.sort( + local_id_to_global_id_by_node_type[supervision_node_type] + ) + ( + sorted_global_node_ids, + local_ids_by_sorted_global_id, + ) = sorted_node_lookup_by_type[supervision_node_type] + output[label_edge_type_to_message_passing_edge_type(label_edge_type)] = ( + _remap_label_tensor_to_local_edge_index( + label_tensor=label_tensor, + sorted_global_node_ids=sorted_global_node_ids, + local_ids_by_sorted_global_id=local_ids_by_sorted_global_id, + ) + ) + return output + + +def remap_labels_to_local_edge_indices( + local_id_to_global_id_by_node_type: dict[NodeType, torch.Tensor], + positive_labels_by_edge_type: dict[EdgeType, torch.Tensor], + negative_labels_by_edge_type: dict[EdgeType, torch.Tensor], +) -> tuple[dict[EdgeType, torch.Tensor], dict[EdgeType, torch.Tensor]]: + """Remap padded global ABLP labels to local label edge indices. + + Positive and negative labels share one sorted node lookup per supervision + node type. Pair order within an anchor is unspecified. + + Args: + local_id_to_global_id_by_node_type: Per node type, a tensor whose + ``i``-th entry is the global id for local node ``i``. + positive_labels_by_edge_type: Per positive-label edge type, a padded + ``[N_anchors, M]`` tensor of global label-node ids. + negative_labels_by_edge_type: Equivalent negative-label tensors. May + be empty. + + Returns: + Positive and negative mappings keyed by message-passing edge type. Each + value is a ``[2, E]`` local label edge index. + """ + sorted_node_lookup_by_type: dict[NodeType, tuple[torch.Tensor, torch.Tensor]] = {} + positive_label_edge_indices = _remap_labels_by_edge_type( + labels_by_edge_type=positive_labels_by_edge_type, + local_id_to_global_id_by_node_type=local_id_to_global_id_by_node_type, + sorted_node_lookup_by_type=sorted_node_lookup_by_type, + ) + negative_label_edge_indices = _remap_labels_by_edge_type( + labels_by_edge_type=negative_labels_by_edge_type, + local_id_to_global_id_by_node_type=local_id_to_global_id_by_node_type, + sorted_node_lookup_by_type=sorted_node_lookup_by_type, + ) + return positive_label_edge_indices, negative_label_edge_indices diff --git a/tests/integration/distributed/graph_store/graph_store_integration_test.py b/tests/integration/distributed/graph_store/graph_store_integration_test.py index bfe41ff7e..9f33a2d04 100644 --- a/tests/integration/distributed/graph_store/graph_store_integration_test.py +++ b/tests/integration/distributed/graph_store/graph_store_integration_test.py @@ -17,6 +17,7 @@ from gigl.common import Uri, UriFactory from gigl.common.logger import Logger from gigl.distributed.dist_ablp_neighborloader import DistABLPLoader +from gigl.distributed.dist_dataset import DistDataset from gigl.distributed.distributed_neighborloader import DistNeighborLoader from gigl.distributed.graph_store.compute import ( init_compute_process, @@ -42,6 +43,14 @@ ) from gigl.utils.data_splitters import DistNodeAnchorLinkSplitter, DistNodeSplitter from gigl.utils.sampling import ABLPInputNodes +from tests.test_assets.distributed.test_dataset import ( + DEFAULT_HETEROGENEOUS_EDGE_INDICES, + STORY, + STORY_TO_USER, + USER, + USER_TO_STORY, + create_heterogeneous_dataset_for_ablp, +) from tests.test_assets.distributed.utils import assert_tensor_equality from tests.test_assets.test_case import DEFAULT_TIMEOUT_SECONDS, TestCase @@ -52,6 +61,38 @@ TEST_NUM_WORKERS = 2 TEST_WORKER_CONCURRENCY = 2 +USER_RATES_STORY = EdgeType(USER, Relation("rates"), STORY) +STORY_RATES_USER = EdgeType(STORY, Relation("rates"), USER) + + +def _build_multi_supervision_ablp_dataset(rank: int, world_size: int) -> DistDataset: + """Build a tiny incoming graph with two objectives and mixed negatives.""" + labels_by_edge_type: dict[ + EdgeType, + tuple[dict[int, list[int]], Optional[dict[int, list[int]]]], + ] = { + USER_TO_STORY: ( + {0: [0, 1], 1: [1, 2], 2: [2, 3], 3: [3], 4: [4]}, + {0: [4], 1: [3], 2: [2], 3: [1], 4: [0]}, + ), + USER_RATES_STORY: ( + {0: [4], 1: [3], 2: [2], 3: [1], 4: [0]}, + None, + ), + } + return create_heterogeneous_dataset_for_ablp( + positive_labels=labels_by_edge_type[USER_TO_STORY][0], + negative_labels=labels_by_edge_type[USER_TO_STORY][1], + train_node_ids=[0, 1, 2], + val_node_ids=[3], + test_node_ids=[4], + edge_indices=DEFAULT_HETEROGENEOUS_EDGE_INDICES, + rank=rank, + world_size=world_size, + edge_dir="in", + labels_by_supervision_edge_type=labels_by_edge_type, + ) + # --------------------------------------------------------------------------- # Tensor helpers @@ -329,6 +370,81 @@ def _run_compute_train_tests( shutdown_compute_process() +def _run_compute_multi_supervision_ablp_test( + client_rank: int, + cluster_info: GraphStoreInfo, + node_type: Optional[NodeType], +) -> None: + """Exercise plural ABLP fetch and collation through real GraphStore RPC.""" + assert node_type == USER + init_compute_process(client_rank, cluster_info, compute_world_backend="gloo") + remote_dataset = RemoteDistDataset( + cluster_info=cluster_info, + local_rank=client_rank, + ) + + inputs = remote_dataset.fetch_ablp_input( + split="train", + rank=torch.distributed.get_rank(), + world_size=torch.distributed.get_world_size(), + anchor_node_type=USER, + supervision_edge_type=[USER_TO_STORY, USER_RATES_STORY], + ) + _assert_ablp_input(cluster_info, inputs) + assert len(inputs) == 1 + server_input = inputs[0] + assert list(server_input.labels) == [STORY_TO_USER, STORY_RATES_USER] + assert_tensor_equality(server_input.anchor_nodes, torch.tensor([0, 1, 2])) + positive_to, negative_to = server_input.labels[STORY_TO_USER] + assert_tensor_equality(positive_to, torch.tensor([[0, 1], [1, 2], [2, 3]]), dim=1) + assert negative_to is not None + assert_tensor_equality(negative_to, torch.tensor([[4], [3], [2]])) + positive_rates, negative_rates = server_input.labels[STORY_RATES_USER] + assert_tensor_equality(positive_rates, torch.tensor([[4], [3], [2]])) + assert negative_rates is None + + loader = DistABLPLoader( + dataset=remote_dataset, + num_neighbors=TEST_NUM_NEIGHBORS, + input_nodes=inputs, + pin_memory_device=TEST_PIN_MEMORY_DEVICE, + num_workers=TEST_NUM_WORKERS, + worker_concurrency=TEST_WORKER_CONCURRENCY, + batch_size=TEST_BATCH_SIZE, + use_edge_index_output=True, + ) + batches = list(loader) + assert len(batches) == 1 + batch = batches[0] + assert isinstance(batch, HeteroData) + assert set(batch.y_positive) == {USER_TO_STORY, USER_RATES_STORY} + assert set(batch.y_negative) == {USER_TO_STORY} + + def _global_pairs( + edge_type: EdgeType, label_edge_index: torch.Tensor + ) -> list[tuple[int, int]]: + anchors = batch[edge_type.src_node_type].node[label_edge_index[0]] + labels = batch[edge_type.dst_node_type].node[label_edge_index[1]] + return sorted(zip(anchors.tolist(), labels.tolist())) + + assert _global_pairs(USER_TO_STORY, batch.y_positive[USER_TO_STORY]) == sorted( + [(0, 0), (0, 1), (1, 1), (1, 2), (2, 2), (2, 3)] + ) + assert _global_pairs(USER_RATES_STORY, batch.y_positive[USER_RATES_STORY]) == [ + (0, 4), + (1, 3), + (2, 2), + ] + assert _global_pairs(USER_TO_STORY, batch.y_negative[USER_TO_STORY]) == [ + (0, 4), + (1, 3), + (2, 2), + ] + + loader.shutdown() + shutdown_compute_process() + + def _run_compute_multiple_loaders_test( client_rank: int, cluster_info: GraphStoreInfo, @@ -718,6 +834,8 @@ class ServerProcessArgs: num_server_sessions: Number of sequential server sessions to run (e.g. one per inference node type). splitter: Optional splitter for node anchor link or node splitting. + dataset_builder: Optional programmatic dataset builder for focused + integration fixtures. """ cluster_info: GraphStoreInfo @@ -726,6 +844,7 @@ class ServerProcessArgs: exception_dict: MutableMapping[str, str] num_server_sessions: int = 1 splitter: Optional[Union[DistNodeAnchorLinkSplitter, DistNodeSplitter]] = None + dataset_builder: Optional[Callable[[int, int], DistDataset]] = None def _run_storage_main_process(args: ServerProcessArgs) -> None: @@ -759,12 +878,15 @@ def _run_storage_main_process(args: ServerProcessArgs) -> None: ) # 2. Build the dataset - dataset = build_storage_dataset( - task_config_uri=args.task_config_uri, - sample_edge_direction=args.sample_edge_direction, - splitter=args.splitter, - tf_record_uri_pattern=".*tfrecord", - ) + if args.dataset_builder is None: + dataset = build_storage_dataset( + task_config_uri=args.task_config_uri, + sample_edge_direction=args.sample_edge_direction, + splitter=args.splitter, + tf_record_uri_pattern=".*tfrecord", + ) + else: + dataset = args.dataset_builder(storage_rank, cluster_info.num_storage_nodes) # 3. Destroy the coordination process group before spawning server # subprocesses. The subprocess will create its own process group on the @@ -890,6 +1012,7 @@ def _launch_graph_store_test( server_splitter: Optional[ Union[DistNodeAnchorLinkSplitter, DistNodeSplitter] ] = None, + server_dataset_builder: Optional[Callable[[int, int], DistDataset]] = None, num_server_sessions: int = 1, ) -> None: """Launch a graph store integration test with the given configuration. @@ -953,6 +1076,7 @@ def _launch_graph_store_test( sample_edge_direction="in", exception_dict=exception_dict, splitter=server_splitter, + dataset_builder=server_dataset_builder, num_server_sessions=num_server_sessions, ) server_process = ctx.Process( @@ -1000,6 +1124,24 @@ def test_homogeneous_training(self): ), ) + def test_multiple_supervision_edge_types_training(self): + """Real RPC preserves two incoming objectives and mixed negatives.""" + task_config_uri = get_mocked_dataset_artifact_metadata()[ + CORA_USER_DEFINED_NODE_ANCHOR_MOCKED_DATASET_INFO.name + ].frozen_gbml_config_uri + cluster_info = self._create_cluster_info( + num_storage_nodes=1, + num_compute_nodes=1, + num_processes_per_compute=1, + ) + self._launch_graph_store_test( + cluster_info=cluster_info, + task_config_uri=task_config_uri, + compute_target=_run_compute_multi_supervision_ablp_test, + node_type=USER, + server_dataset_builder=_build_multi_supervision_ablp_dataset, + ) + def test_multiple_loaders_in_graph_store(self): """Test that multiple loader instances (2 ABLP + 2 DistNeighborLoader) can work in parallel, followed by another (ABLP, DistNeighborLoader) pair sequentially. diff --git a/tests/test_assets/distributed/test_dataset.py b/tests/test_assets/distributed/test_dataset.py index eb0b852c5..6ea34184f 100644 --- a/tests/test_assets/distributed/test_dataset.py +++ b/tests/test_assets/distributed/test_dataset.py @@ -247,6 +247,12 @@ def create_heterogeneous_dataset_for_ablp( rank: int = 0, world_size: int = 1, edge_dir: Literal["in", "out"] = "out", + labels_by_supervision_edge_type: Optional[ + dict[ + EdgeType, + tuple[dict[int, list[int]], Optional[dict[int, list[int]]]], + ] + ] = None, ) -> DistDataset: """Create a heterogeneous test dataset for ABLP with label edges and train/val/test splits. @@ -278,6 +284,11 @@ def create_heterogeneous_dataset_for_ablp( rank: Rank of the current process. Defaults to 0. world_size: Total number of processes. Defaults to 1. edge_dir: Edge direction ("in" or "out"). Defaults to "out". + labels_by_supervision_edge_type: Optional plural label schema keyed by + canonical anchor-outward supervision edge type. When provided, + ``positive_labels``, ``negative_labels``, and + ``supervision_edge_type`` supply only the backward-compatible scalar + defaults and this mapping is authoritative. Returns: A DistDataset instance with the specified configuration and splits. @@ -295,48 +306,29 @@ def create_heterogeneous_dataset_for_ablp( ... edge_indices=DEFAULT_HETEROGENEOUS_EDGE_INDICES, ... ) """ - # Set default supervision edge type + # Set default supervision edge type and normalize the scalar convenience + # form to the plural schema used below. if supervision_edge_type is None: supervision_edge_type = EdgeType(src_node_type, Relation("to"), dst_node_type) + if labels_by_supervision_edge_type is None: + labels_by_supervision_edge_type = { + supervision_edge_type: (positive_labels, negative_labels) + } - # Validate that all split node IDs have positive labels + # Validate that all split node IDs have positive labels for every objective. all_split_node_ids = set(train_node_ids) | set(val_node_ids) | set(test_node_ids) - missing_nodes = all_split_node_ids - set(positive_labels.keys()) - if missing_nodes: - raise ValueError( - f"Node IDs {missing_nodes} are in train/val/test splits but not in positive_labels" - ) - - # When edge_dir="in", the DistNodeAnchorLinkSplitter reverses the label edge type - # (supervision_edge_type is always provided outward, e.g. USER_TO_STORY), so the splitter - # looks for message_passing_to_positive_label(reverse(supervision_edge_type)) in the graph. - # We must store the label edges under that reversed key for the splitter and ABLP loader to - # find them. For edge_dir="out" no reversal is needed. - actual_label_supervision_edge_type = ( - reverse_edge_type(supervision_edge_type) - if edge_dir == "in" - else supervision_edge_type - ) - positive_label_edge_type = message_passing_to_positive_label( - actual_label_supervision_edge_type - ) - negative_label_edge_type = message_passing_to_negative_label( - actual_label_supervision_edge_type - ) - - # Convert positive_labels dict to COO edge index - pos_src, pos_dst = [], [] - for node_id, dst_ids in positive_labels.items(): - for dst_id in dst_ids: - pos_src.append(node_id) - pos_dst.append(dst_id) - # For edge_dir="in", the reversed label edge type is (STORY, ..., USER), so - # row 0 = story IDs, row 1 = user IDs. For edge_dir="out" it's the natural direction. - positive_label_edge_index = ( - torch.tensor([pos_dst, pos_src]) - if edge_dir == "in" - else torch.tensor([pos_src, pos_dst]) - ) + for edge_type, (edge_positive_labels, _) in labels_by_supervision_edge_type.items(): + if edge_type.src_node_type != src_node_type: + raise ValueError( + f"Supervision edge type {edge_type} must use shared anchor " + f"node type {src_node_type} as its source." + ) + missing_nodes = all_split_node_ids - set(edge_positive_labels.keys()) + if missing_nodes: + raise ValueError( + f"Node IDs {missing_nodes} are in train/val/test splits but not " + f"in positive labels for {edge_type}." + ) # Derive node counts from edge indices by collecting max node ID per node type node_counts: dict[NodeType, int] = {} @@ -347,53 +339,105 @@ def create_heterogeneous_dataset_for_ablp( node_counts[src_type] = int(max(node_counts.get(src_type, 0), src_max)) node_counts[dst_type] = int(max(node_counts.get(dst_type, 0), dst_max)) - # Also account for nodes in positive labels - node_counts[src_node_type] = max( - node_counts.get(src_node_type, 0), max(positive_labels.keys()) + 1 - ) - node_counts[dst_node_type] = max( - node_counts.get(dst_node_type, 0), - max(max(stories) for stories in positive_labels.values()) + 1, - ) - # Set up edge partition books and edge indices edge_partition_book = { edge_type: torch.zeros(edge_index.shape[1], dtype=torch.int64) for edge_type, edge_index in edge_indices.items() } - edge_partition_book[positive_label_edge_type] = torch.zeros( - len(pos_src), dtype=torch.int64 - ) partitioned_edge_index = { edge_type: GraphPartitionData(edge_index=edge_index, edge_ids=None) for edge_type, edge_index in edge_indices.items() } - partitioned_edge_index[positive_label_edge_type] = GraphPartitionData( - edge_index=positive_label_edge_index, - edge_ids=None, - ) - if negative_labels is not None: - # Convert negative_labels dict to COO edge index - neg_src, neg_dst = [], [] - for node_id, dst_ids in negative_labels.items(): - for dst_id in dst_ids: - neg_src.append(node_id) - neg_dst.append(dst_id) - negative_label_edge_index = ( - torch.tensor([neg_dst, neg_src]) + for outward_edge_type, ( + edge_positive_labels, + edge_negative_labels, + ) in labels_by_supervision_edge_type.items(): + registered_edge_type = ( + reverse_edge_type(outward_edge_type) if edge_dir == "in" - else torch.tensor([neg_src, neg_dst]) + else outward_edge_type ) - edge_partition_book[negative_label_edge_type] = torch.zeros( - len(neg_src), dtype=torch.int64 + positive_label_edge_type = message_passing_to_positive_label( + registered_edge_type ) - partitioned_edge_index[negative_label_edge_type] = GraphPartitionData( - edge_index=negative_label_edge_index, + positive_src = [ + anchor + for anchor, destination_ids in edge_positive_labels.items() + for _ in destination_ids + ] + positive_dst = [ + destination + for destination_ids in edge_positive_labels.values() + for destination in destination_ids + ] + positive_label_edge_index = ( + torch.tensor([positive_dst, positive_src]) + if edge_dir == "in" + else torch.tensor([positive_src, positive_dst]) + ) + edge_partition_book[positive_label_edge_type] = torch.zeros( + len(positive_src), dtype=torch.int64 + ) + partitioned_edge_index[positive_label_edge_type] = GraphPartitionData( + edge_index=positive_label_edge_index, edge_ids=None, ) + node_counts[src_node_type] = max( + node_counts.get(src_node_type, 0), max(edge_positive_labels) + 1 + ) + destination_node_type = outward_edge_type.dst_node_type + node_counts[destination_node_type] = max( + node_counts.get(destination_node_type, 0), + max( + destination + for destinations in edge_positive_labels.values() + for destination in destinations + ) + + 1, + ) + + if edge_negative_labels is not None: + negative_label_edge_type = message_passing_to_negative_label( + registered_edge_type + ) + negative_src = [ + anchor + for anchor, destination_ids in edge_negative_labels.items() + for _ in destination_ids + ] + negative_dst = [ + destination + for destination_ids in edge_negative_labels.values() + for destination in destination_ids + ] + negative_label_edge_index = ( + torch.tensor([negative_dst, negative_src]) + if edge_dir == "in" + else torch.tensor([negative_src, negative_dst]) + ) + edge_partition_book[negative_label_edge_type] = torch.zeros( + len(negative_src), dtype=torch.int64 + ) + partitioned_edge_index[negative_label_edge_type] = GraphPartitionData( + edge_index=negative_label_edge_index, + edge_ids=None, + ) + node_counts[src_node_type] = max( + node_counts.get(src_node_type, 0), max(edge_negative_labels) + 1 + ) + node_counts[destination_node_type] = max( + node_counts.get(destination_node_type, 0), + max( + destination + for destinations in edge_negative_labels.values() + for destination in destinations + ) + + 1, + ) + # Partition books filled with zeros assign all nodes to partition 0 node_partition_book = { node_type: torch.zeros(count, dtype=torch.int64) @@ -424,7 +468,8 @@ def create_heterogeneous_dataset_for_ablp( # Calculate split ratios based on provided node IDs. # With identity hash (x + 1), nodes are split by their ID values: # - Lower IDs -> train, middle IDs -> val, higher IDs -> test - total_nodes = len(positive_labels) + first_positive_labels = next(iter(labels_by_supervision_edge_type.values()))[0] + total_nodes = len(first_positive_labels) num_val = len(val_node_ids) / total_nodes num_test = len(test_node_ids) / total_nodes @@ -439,7 +484,7 @@ def _identity_hash(x: torch.Tensor) -> torch.Tensor: num_val=num_val, num_test=num_test, hash_function=_identity_hash, - supervision_edge_types=[supervision_edge_type], + supervision_edge_types=list(labels_by_supervision_edge_type), should_convert_labels_to_edges=True, ) diff --git a/tests/unit/distributed/dist_ablp_neighborloader_test.py b/tests/unit/distributed/dist_ablp_neighborloader_test.py index 31d3d1cbc..36f57e04d 100644 --- a/tests/unit/distributed/dist_ablp_neighborloader_test.py +++ b/tests/unit/distributed/dist_ablp_neighborloader_test.py @@ -11,7 +11,10 @@ from torch_geometric.data import Data, HeteroData from gigl.distributed.dataset_factory import build_dataset -from gigl.distributed.dist_ablp_neighborloader import DistABLPLoader +from gigl.distributed.dist_ablp_neighborloader import ( + DistABLPLoader, + _convert_graph_store_ablp_inputs, +) from gigl.distributed.dist_dataset import DistDataset from gigl.distributed.dist_partitioner import DistPartitioner from gigl.distributed.dist_range_partitioner import DistRangePartitioner @@ -36,6 +39,7 @@ to_homogeneous, ) from gigl.utils.data_splitters import DistNodeAnchorLinkSplitter +from gigl.utils.sampling import ABLPInputNodes from tests.test_assets.distributed.utils import ( assert_tensor_equality, create_test_process_group, @@ -358,12 +362,13 @@ def _run_distributed_ablp_neighbor_loader_multiple_supervision_edge_types( expected=expected_positive_labels[edge_type], ) if expected_negative_labels is not None: - _assert_labels( - anchor_nodes=datum[edge_type[anchor_index]].node, - supervision_nodes=datum[edge_type[supervision_index]].node, - y=datum.y_negative[edge_type], - expected=expected_negative_labels[edge_type], - ) + for edge_type, expected_labels in expected_negative_labels.items(): + _assert_labels( + anchor_nodes=datum[edge_type[anchor_index]].node, + supervision_nodes=datum[edge_type[supervision_index]].node, + y=datum.y_negative[edge_type], + expected=expected_labels, + ) else: assert not hasattr(datum, "y_negative") @@ -416,6 +421,177 @@ def _run_distributed_ablp_neighbor_loader_multiple_supervision_edge_types( shutdown_rpc() +def _global_pair_set( + anchor_node: torch.Tensor, + label_node: torch.Tensor, + label_dict: dict[int, torch.Tensor], +) -> list[tuple[int, int]]: + """Convert a label dictionary to sorted global-id pairs. + + Anchors and labels are looked up in separate node maps so this works for + heterogeneous graphs, where the two live in different node stores. Pass the + same tensor twice for a homogeneous graph. + """ + pairs: list[tuple[int, int]] = [] + for local_anchor, local_labels in label_dict.items(): + global_anchor = int(anchor_node[local_anchor].item()) + for local_label in local_labels.tolist(): + pairs.append((global_anchor, int(label_node[local_label].item()))) + return sorted(pairs) + + +def _global_pair_set_from_edge_index( + anchor_node: torch.Tensor, + label_node: torch.Tensor, + label_edge_index: torch.Tensor, +) -> list[tuple[int, int]]: + """Convert a label edge index to sorted global-id pairs. + + Takes separate anchor and label node maps for the same reason as + :func:`_global_pair_set`. + """ + return sorted( + ( + int(anchor_node[local_anchor].item()), + int(label_node[local_label].item()), + ) + for local_anchor, local_label in label_edge_index.t().tolist() + ) + + +def _collect_homogeneous_labels( + _, + return_dict, + use_edge_index_output: bool, + dataset: DistDataset, + input_nodes: torch.Tensor, + batch_size: int, + has_negatives: bool, +): + """Run one loader format and return its sorted global-id label pairs.""" + create_test_process_group() + loader = DistABLPLoader( + dataset=dataset, + num_neighbors=[2, 2], + input_nodes=input_nodes, + batch_size=batch_size, + pin_memory_device=torch.device("cpu"), + use_edge_index_output=use_edge_index_output, + ) + positive_pairs: list[tuple[int, int]] = [] + negative_pairs: list[tuple[int, int]] = [] + for datum in loader: + assert isinstance(datum, Data) + node = datum.node + if use_edge_index_output: + assert isinstance(datum.y_positive, torch.Tensor) + assert datum.y_positive.size(0) == 2 + positive_pairs.extend( + _global_pair_set_from_edge_index(node, node, datum.y_positive) + ) + else: + positive_pairs.extend(_global_pair_set(node, node, datum.y_positive)) + if has_negatives: + if use_edge_index_output: + assert isinstance(datum.y_negative, torch.Tensor) + assert datum.y_negative.size(0) == 2 + negative_pairs.extend( + _global_pair_set_from_edge_index(node, node, datum.y_negative) + ) + else: + negative_pairs.extend(_global_pair_set(node, node, datum.y_negative)) + else: + assert not hasattr(datum, "y_negative"), ( + f"expected no negatives, got {getattr(datum, 'y_negative', None)}" + ) + return_dict[use_edge_index_output] = ( + sorted(positive_pairs), + sorted(negative_pairs), + ) + shutdown_rpc() + + +def _edge_type_key(edge_type: EdgeType) -> tuple[str, ...]: + """Canonical dict key for an edge type. + + Label edge types reach collation as plain tuples on some paths and as + ``EdgeType`` on others, and the two stringify differently (``"('a', 'to', + 'b')"`` vs ``'a-to-b'``). Normalizing to plain strings keeps keys comparable + across the process boundary regardless of which arrives. + """ + return tuple(str(part) for part in edge_type) + + +def _accumulate_heterogeneous_pairs( + data: HeteroData, + labels_by_edge_type: dict[EdgeType, Union[torch.Tensor, dict[int, torch.Tensor]]], + use_edge_index_output: bool, + into: dict[tuple[str, ...], list[tuple[int, int]]], +) -> None: + """Accumulate one batch's labels as global (anchor, label) id pairs. + + Anchors and supervision nodes are resolved through the node maps of the edge + type's src and dst stores respectively, so a label remapped against the wrong + node type surfaces as a wrong global id rather than passing silently. + """ + anchor_index = 0 + supervision_index = 2 + for edge_type, labels in labels_by_edge_type.items(): + anchor_node = data[edge_type[anchor_index]].node + label_node = data[edge_type[supervision_index]].node + if use_edge_index_output: + assert isinstance(labels, torch.Tensor), f"{edge_type}: {type(labels)}" + assert labels.size(0) == 2 + pairs = _global_pair_set_from_edge_index(anchor_node, label_node, labels) + else: + assert not isinstance(labels, torch.Tensor), f"{edge_type}: {type(labels)}" + pairs = _global_pair_set(anchor_node, label_node, labels) + into[_edge_type_key(edge_type)].extend(pairs) + + +def _collect_heterogeneous_labels( + _, + return_dict, + use_edge_index_output: bool, + dataset: DistDataset, + input_nodes: tuple[NodeType, torch.Tensor], + supervision_edge_types: list[EdgeType], + batch_size: int, +): + """Run one loader format on a heterogeneous graph and return its label pairs. + + Results are keyed by ``str(edge_type)`` so they cross the process boundary as + plain data. + """ + create_test_process_group() + loader = DistABLPLoader( + dataset=dataset, + num_neighbors=[2, 2], + input_nodes=input_nodes, + batch_size=batch_size, + pin_memory_device=torch.device("cpu"), + supervision_edge_type=supervision_edge_types, + use_edge_index_output=use_edge_index_output, + ) + positive_pairs: dict[tuple[str, ...], list[tuple[int, int]]] = defaultdict(list) + negative_pairs: dict[tuple[str, ...], list[tuple[int, int]]] = defaultdict(list) + for datum in loader: + assert isinstance(datum, HeteroData) + # Several supervision edge types keep y_positive / y_negative in their + # dict-of-edge-type form rather than collapsing to a bare value. + _accumulate_heterogeneous_pairs( + datum, datum.y_positive, use_edge_index_output, positive_pairs + ) + _accumulate_heterogeneous_pairs( + datum, datum.y_negative, use_edge_index_output, negative_pairs + ) + return_dict[use_edge_index_output] = ( + {edge_type: sorted(pairs) for edge_type, pairs in positive_pairs.items()}, + {edge_type: sorted(pairs) for edge_type, pairs in negative_pairs.items()}, + ) + shutdown_rpc() + + class DistABLPLoaderTest(TestCase): def tearDown(self): if torch.distributed.is_initialized(): @@ -425,6 +601,154 @@ def tearDown(self): torch.distributed.destroy_process_group() super().tearDown() + def test_convert_graph_store_inputs_with_multiple_supervision_types(self): + """GraphStore conversion retains plural and per-type label topology.""" + positive_a_to_b = message_passing_to_positive_label(_A_TO_B) + negative_a_to_b = message_passing_to_negative_label(_A_TO_B) + positive_a_to_c = message_passing_to_positive_label(_A_TO_C) + anchors = torch.tensor([10, 11]) + input_nodes = { + 0: ABLPInputNodes( + anchor_nodes=anchors, + anchor_node_type=_A, + labels={ + _A_TO_B: ( + torch.tensor([[12], [13]]), + torch.tensor([[14], [15]]), + ), + _A_TO_C: (torch.tensor([[20], [21]]), None), + }, + ) + } + + ( + sampler_inputs, + input_type, + supervision_edge_types, + positive_edge_types, + negative_edge_types, + ) = _convert_graph_store_ablp_inputs( + input_nodes=input_nodes, + num_storage_nodes=2, + edge_types=[ + positive_a_to_b, + negative_a_to_b, + positive_a_to_c, + ], + edge_dir="out", + ) + + self.assertEqual(input_type, _A) + self.assertEqual(supervision_edge_types, [_A_TO_B, _A_TO_C]) + self.assertEqual(positive_edge_types, [positive_a_to_b, positive_a_to_c]) + self.assertEqual(negative_edge_types, [negative_a_to_b]) + self.assertEqual(len(sampler_inputs), 2) + self.assert_tensor_equality(sampler_inputs[0].node, anchors) + self.assertEqual( + list(sampler_inputs[0].positive_label_by_edge_types), + [positive_a_to_b, positive_a_to_c], + ) + self.assertEqual( + list(sampler_inputs[0].negative_label_by_edge_types), + [negative_a_to_b], + ) + self.assertEqual(sampler_inputs[1].node.numel(), 0) + self.assertEqual( + tuple( + sampler_inputs[1].positive_label_by_edge_types[positive_a_to_c].shape + ), + (0, 0), + ) + self.assertEqual( + tuple( + sampler_inputs[1].negative_label_by_edge_types[negative_a_to_b].shape + ), + (0, 0), + ) + + def test_convert_graph_store_inputs_rejects_malformed_schema(self): + """Every server entry must agree with topology and tensor schema.""" + positive_a_to_b = message_passing_to_positive_label(_A_TO_B) + negative_a_to_b = message_passing_to_negative_label(_A_TO_B) + valid = ABLPInputNodes( + anchor_nodes=torch.tensor([10]), + anchor_node_type=_A, + labels={ + _A_TO_B: ( + torch.tensor([[12]]), + torch.tensor([[13]]), + ) + }, + ) + cases = { + "out-of-range rank": { + 1: valid, + }, + "wrong anchor type": { + 0: ABLPInputNodes( + anchor_nodes=torch.tensor([10]), + anchor_node_type=_B, + labels=valid.labels, + ) + }, + "non-integral anchors": { + 0: ABLPInputNodes( + anchor_nodes=torch.tensor([10.0]), + anchor_node_type=_A, + labels=valid.labels, + ) + }, + "wrong positive rows": { + 0: ABLPInputNodes( + anchor_nodes=torch.tensor([10]), + anchor_node_type=_A, + labels={ + _A_TO_B: ( + torch.tensor([[12], [14]]), + torch.tensor([[13]]), + ) + }, + ) + }, + "missing topology negative": { + 0: ABLPInputNodes( + anchor_nodes=torch.tensor([10]), + anchor_node_type=_A, + labels={_A_TO_B: (torch.tensor([[12]]), None)}, + ) + }, + } + for name, input_nodes in cases.items(): + with self.subTest(name=name): + with self.assertRaises(ValueError): + _convert_graph_store_ablp_inputs( + input_nodes=input_nodes, + num_storage_nodes=1, + edge_types=[positive_a_to_b, negative_a_to_b], + edge_dir="out", + ) + with self.assertRaises(ValueError): + _convert_graph_store_ablp_inputs( + input_nodes={0: valid}, + num_storage_nodes=1, + edge_types=[positive_a_to_b], + edge_dir="out", + ) + with self.assertRaises(ValueError): + _convert_graph_store_ablp_inputs( + input_nodes={ + 0: valid, + 1: ABLPInputNodes( + anchor_nodes=torch.tensor([11]), + anchor_node_type=_A, + labels={_A_TO_C: (torch.tensor([[20]]), None)}, + ), + }, + num_storage_nodes=2, + edge_types=[positive_a_to_b, negative_a_to_b], + edge_dir="out", + ) + @parameterized.expand( [ param( @@ -556,6 +880,216 @@ def test_ablp_dataloader( ), ) + @parameterized.expand( + [ + param( + "positive and negative", + labeled_edges={ + _POSITIVE_EDGE_TYPE: torch.tensor([[10, 15], [15, 16]]), + _NEGATIVE_EDGE_TYPE: torch.tensor( + [[10, 10, 11, 15], [13, 16, 14, 17]] + ), + }, + input_nodes=torch.tensor([10, 15]), + batch_size=2, + has_negatives=True, + ), + param( + "positive only", + labeled_edges={_POSITIVE_EDGE_TYPE: torch.tensor([[10, 15], [15, 16]])}, + input_nodes=torch.tensor([10, 15]), + batch_size=2, + has_negatives=False, + ), + # Anchor 11 has message-passing edges (11 -> {13, 17}) but is the + # source of NO positive-label edge, so its positive-label row is + # all-padding and y_positive[11] is a guaranteed-empty tensor. This + # exercises the empty-anchor branch end-to-end for both outputs. + param( + "guaranteed empty positive anchor", + labeled_edges={ + _POSITIVE_EDGE_TYPE: torch.tensor([[10, 15], [15, 16]]), + _NEGATIVE_EDGE_TYPE: torch.tensor( + [[10, 10, 11, 15], [13, 16, 14, 17]] + ), + }, + input_nodes=torch.tensor([10, 11, 15]), + batch_size=3, + has_negatives=True, + ), + ] + ) + def test_edge_index_output_matches_dict_output( + self, _, labeled_edges, input_nodes, batch_size, has_negatives + ): + """Both output formats contain the same global anchor-label pairs.""" + edge_index = { + DEFAULT_HOMOGENEOUS_EDGE_TYPE: torch.tensor( + [[10, 10, 11, 11, 15, 15, 16, 16], [11, 12, 13, 17, 13, 14, 12, 14]] + ), + } + edge_index.update(labeled_edges) + partition_output = PartitionOutput( + node_partition_book=to_heterogeneous_node(torch.zeros(18)), + edge_partition_book={ + e_type: torch.zeros(int(e_idx.max().item() + 1)) + for e_type, e_idx in edge_index.items() + }, + partitioned_edge_index={ + etype: GraphPartitionData( + edge_index=idx, edge_ids=torch.arange(idx.size(1)) + ) + for etype, idx in edge_index.items() + }, + partitioned_edge_features=None, + partitioned_node_features=None, + partitioned_negative_labels=None, + partitioned_positive_labels=None, + partitioned_node_labels=None, + ) + dataset = DistDataset(rank=0, world_size=1, edge_dir="out") + dataset.build(partition_output=partition_output) + + manager = mp.Manager() + return_dict = manager.dict() + for use_edge_index_output in (False, True): + mp.spawn( + fn=_collect_homogeneous_labels, + args=( + return_dict, + use_edge_index_output, + dataset, + input_nodes, + batch_size, + has_negatives, + ), + ) + self.assertEqual(return_dict[False][0], return_dict[True][0]) + self.assertEqual(return_dict[False][1], return_dict[True][1]) + + @parameterized.expand( + [ + param( + "edge_dir=out", + edge_dir="out", + edge_index={ + _A_TO_B: torch.tensor([[10, 10], [11, 12]]), + message_passing_to_positive_label(_A_TO_B): torch.tensor( + [[10, 10], [13, 14]] + ), + message_passing_to_negative_label(_A_TO_B): torch.tensor( + [[10, 10], [15, 16]] + ), + _A_TO_C: torch.tensor([[10, 10], [20, 21]]), + message_passing_to_positive_label(_A_TO_C): torch.tensor( + [[10, 10], [22, 23]] + ), + message_passing_to_negative_label(_A_TO_C): torch.tensor( + [[10, 10], [24, 25]] + ), + }, + ), + # edge_dir="in" stores the supervision edge types reversed, so their + # dst node type is the anchor type rather than the supervision type, + # while the label edge types reaching collation stay outward-facing. + # A loader that resolves supervision nodes off the former instead of + # the latter fails here and passes the edge_dir="out" case above. + param( + "edge_dir=in", + edge_dir="in", + edge_index={ + _B_TO_A: torch.tensor([[11, 12], [10, 10]]), + message_passing_to_positive_label(_B_TO_A): torch.tensor( + [[13, 14], [10, 10]] + ), + message_passing_to_negative_label(_B_TO_A): torch.tensor( + [[15, 16], [10, 10]] + ), + _C_TO_A: torch.tensor([[20, 21], [10, 10]]), + message_passing_to_positive_label(_C_TO_A): torch.tensor( + [[22, 23], [10, 10]] + ), + message_passing_to_negative_label(_C_TO_A): torch.tensor( + [[24, 25], [10, 10]] + ), + }, + ), + ] + ) + def test_heterogeneous_edge_index_output_matches_dict_output( + self, + _, + edge_dir: Literal["in", "out"], + edge_index: dict[EdgeType, torch.Tensor], + ): + """Both output formats agree on a heterogeneous graph, for either edge_dir. + + Anchors are node type ``a``; supervision nodes are ``b`` and ``c``. Because + the anchor and supervision node maps are distinct, a label remapped against + the wrong node type yields a wrong global id here instead of going unnoticed. + """ + nodes: dict[NodeType, list[torch.Tensor]] = defaultdict(list) + for edge_type, edge_idx in edge_index.items(): + nodes[edge_type[0]].append(edge_idx[0]) + nodes[edge_type[2]].append(edge_idx[1]) + partition_output = PartitionOutput( + node_partition_book={ + node_type: torch.zeros(int(torch.cat(node_ids).max().item() + 1)) + for node_type, node_ids in nodes.items() + }, + edge_partition_book={ + e_type: torch.zeros(int(e_idx.max().item() + 1)) + for e_type, e_idx in edge_index.items() + }, + partitioned_edge_index={ + etype: GraphPartitionData( + edge_index=idx, edge_ids=torch.arange(idx.size(1)) + ) + for etype, idx in edge_index.items() + }, + partitioned_edge_features=None, + partitioned_node_features=None, + partitioned_negative_labels=None, + partitioned_positive_labels=None, + partitioned_node_labels=None, + ) + dataset = DistDataset(rank=0, world_size=1, edge_dir=edge_dir) + dataset.build(partition_output=partition_output) + + manager = mp.Manager() + return_dict = manager.dict() + for use_edge_index_output in (False, True): + mp.spawn( + fn=_collect_heterogeneous_labels, + args=( + return_dict, + use_edge_index_output, + dataset, + (_A, torch.tensor([10])), + [_A_TO_B, _A_TO_C], + 1, # batch_size + ), + ) + # Both formats resolve to the same global ids... + self.assertEqual(dict(return_dict[False][0]), dict(return_dict[True][0])) + self.assertEqual(dict(return_dict[False][1]), dict(return_dict[True][1])) + # ...and did so over labels that actually exist, so equality above cannot + # hold vacuously. + self.assertEqual( + dict(return_dict[True][0]), + { + _edge_type_key(_A_TO_B): [(10, 13), (10, 14)], + _edge_type_key(_A_TO_C): [(10, 22), (10, 23)], + }, + ) + self.assertEqual( + dict(return_dict[True][1]), + { + _edge_type_key(_A_TO_B): [(10, 15), (10, 16)], + _edge_type_key(_A_TO_C): [(10, 24), (10, 25)], + }, + ) + def test_cora_supervised(self): create_test_process_group() cora_supervised_info = get_mocked_dataset_artifact_metadata()[ @@ -811,6 +1345,45 @@ def test_toy_heterogeneous_ablp( _A_TO_C: {10: torch.tensor([24, 25])}, }, ), + param( + "negative edges on only one supervision type", + edge_dir="out", + edge_index={ + _A_TO_B: torch.tensor([[10, 10], [11, 12]]), + message_passing_to_positive_label(_A_TO_B): torch.tensor( + [[10, 10], [13, 14]] + ), + message_passing_to_negative_label(_A_TO_B): torch.tensor( + [[10, 10], [15, 16]] + ), + _A_TO_C: torch.tensor([[10, 10], [20, 21]]), + message_passing_to_positive_label(_A_TO_C): torch.tensor( + [[10, 10], [22, 23]] + ), + }, + supervision_edge_types=[_A_TO_B, _A_TO_C], + expected_node={ + _A: torch.tensor([10]), + _B: torch.tensor([11, 12, 13, 14, 15, 16]), + _C: torch.tensor([20, 21, 22, 23]), + }, + expected_batch={ + _A: torch.tensor([10]), + _B: None, + _C: None, + }, + expected_edges={ + _A_TO_B: (torch.tensor([10, 10]), torch.tensor([11, 12])), + _A_TO_C: (torch.tensor([10, 10]), torch.tensor([20, 21])), + }, + expected_positive_labels={ + _A_TO_B: {10: torch.tensor([13, 14])}, + _A_TO_C: {10: torch.tensor([22, 23])}, + }, + expected_negative_labels={ + _A_TO_B: {10: torch.tensor([15, 16])}, + }, + ), # https://is.gd/mO5cpW param( "same nodes, different relation", diff --git a/tests/unit/distributed/dist_server_test.py b/tests/unit/distributed/dist_server_test.py index cfa4f536b..c4d83a7b8 100644 --- a/tests/unit/distributed/dist_server_test.py +++ b/tests/unit/distributed/dist_server_test.py @@ -13,7 +13,7 @@ RegisterBackendRequest, ) from gigl.distributed.graph_store.sharding import ServerSlice -from gigl.src.common.types.graph_data import Relation +from gigl.src.common.types.graph_data import EdgeType, Relation from tests.test_assets.distributed.test_dataset import ( DEFAULT_HETEROGENEOUS_EDGE_INDICES, DEFAULT_HOMOGENEOUS_EDGE_INDEX, @@ -380,13 +380,15 @@ def test_get_ablp_input(self) -> None: for split, expected_user_ids in split_to_user_ids.items(): with self.subTest(split=split): - anchor_nodes, pos_labels, neg_labels = server.get_ablp_input( + ablp_input = server.get_ablp_input( FetchABLPInputRequest( split=split, node_type=USER, - supervision_edge_type=USER_TO_STORY, + supervision_edge_types=(USER_TO_STORY,), ) ) + anchor_nodes = ablp_input.anchor_nodes + pos_labels, neg_labels = ablp_input.labels[USER_TO_STORY] # Verify anchor nodes match expected users self.assert_tensor_equality( @@ -404,6 +406,70 @@ def test_get_ablp_input(self) -> None: assert neg_labels is not None self.assert_tensor_equality(neg_labels, torch.tensor(expected_negative)) + def test_get_ablp_input_with_multiple_supervision_edge_types(self) -> None: + """One server request returns all label types with per-type negatives.""" + create_test_process_group() + user_rates_story = EdgeType(USER, Relation("rates"), STORY) + labels_by_edge_type: dict[ + EdgeType, + tuple[dict[int, list[int]], dict[int, list[int]] | None], + ] = { + USER_TO_STORY: ( + {0: [0, 1], 1: [1, 2], 2: [2, 3], 3: [3], 4: [4]}, + {0: [4], 1: [3], 2: [2], 3: [1], 4: [0]}, + ), + user_rates_story: ( + {0: [4], 1: [3], 2: [2], 3: [1], 4: [0]}, + None, + ), + } + dataset = create_heterogeneous_dataset_for_ablp( + positive_labels=labels_by_edge_type[USER_TO_STORY][0], + negative_labels=labels_by_edge_type[USER_TO_STORY][1], + train_node_ids=[0, 1, 2], + val_node_ids=[3], + test_node_ids=[4], + edge_indices=DEFAULT_HETEROGENEOUS_EDGE_INDICES, + labels_by_supervision_edge_type=labels_by_edge_type, + ) + server = dist_server.DistServer(dataset) + + result = server.get_ablp_input( + FetchABLPInputRequest( + split="train", + node_type=USER, + supervision_edge_types=(USER_TO_STORY, user_rates_story), + ) + ) + + self.assert_tensor_equality(result.anchor_nodes, torch.tensor([0, 1, 2])) + self.assertEqual(list(result.labels), [USER_TO_STORY, user_rates_story]) + positive_to, negative_to = result.labels[USER_TO_STORY] + self.assert_tensor_equality( + positive_to, torch.tensor([[0, 1], [1, 2], [2, 3]]), dim=1 + ) + assert negative_to is not None + self.assert_tensor_equality(negative_to, torch.tensor([[4], [3], [2]])) + positive_rates, negative_rates = result.labels[user_rates_story] + self.assert_tensor_equality(positive_rates, torch.tensor([[4], [3], [2]])) + self.assertIsNone(negative_rates) + + for supervision_edge_types in ( + (), + (USER_TO_STORY, USER_TO_STORY), + (EdgeType(STORY, Relation("to"), USER),), + (EdgeType(USER, Relation("missing"), STORY),), + ): + with self.subTest(supervision_edge_types=supervision_edge_types): + with self.assertRaises(ValueError): + server.get_ablp_input( + FetchABLPInputRequest( + split="train", + node_type=USER, + supervision_edge_types=supervision_edge_types, + ) + ) + def test_get_ablp_input_with_server_slicing(self) -> None: """Test get_ablp_input with server slices to verify sharding.""" create_test_process_group() @@ -434,11 +500,11 @@ def test_get_ablp_input_with_server_slicing(self) -> None: server = dist_server.DistServer(dataset) # Get training input for first half - anchor_nodes_0, pos_labels_0, neg_labels_0 = server.get_ablp_input( + ablp_input_0 = server.get_ablp_input( FetchABLPInputRequest( split="train", node_type=USER, - supervision_edge_type=USER_TO_STORY, + supervision_edge_types=(USER_TO_STORY,), server_slice=ServerSlice( server_rank=0, start_numerator=0, end_numerator=1, denominator=2 ), @@ -446,16 +512,20 @@ def test_get_ablp_input_with_server_slicing(self) -> None: ) # Get training input for second half - anchor_nodes_1, pos_labels_1, neg_labels_1 = server.get_ablp_input( + ablp_input_1 = server.get_ablp_input( FetchABLPInputRequest( split="train", node_type=USER, - supervision_edge_type=USER_TO_STORY, + supervision_edge_types=(USER_TO_STORY,), server_slice=ServerSlice( server_rank=0, start_numerator=1, end_numerator=2, denominator=2 ), ) ) + anchor_nodes_0 = ablp_input_0.anchor_nodes + pos_labels_0, neg_labels_0 = ablp_input_0.labels[USER_TO_STORY] + anchor_nodes_1 = ablp_input_1.anchor_nodes + pos_labels_1, neg_labels_1 = ablp_input_1.labels[USER_TO_STORY] # Train nodes [0, 1, 2, 3] should be split across ranks rank_0_user_ids = [0, 1] @@ -502,7 +572,7 @@ def test_get_ablp_input_invalid_split(self) -> None: FetchABLPInputRequest( split="invalid", node_type=USER, - supervision_edge_type=USER_TO_STORY, + supervision_edge_types=(USER_TO_STORY,), ) ) @@ -529,13 +599,15 @@ def test_get_ablp_input_without_negative_labels(self) -> None: ) server = dist_server.DistServer(dataset) - anchor_nodes, pos_labels, neg_labels = server.get_ablp_input( + ablp_input = server.get_ablp_input( FetchABLPInputRequest( split="train", node_type=USER, - supervision_edge_type=USER_TO_STORY, + supervision_edge_types=(USER_TO_STORY,), ) ) + anchor_nodes = ablp_input.anchor_nodes + pos_labels, neg_labels = ablp_input.labels[USER_TO_STORY] # Verify train split returns the expected users self.assert_tensor_equality(anchor_nodes, torch.tensor(train_user_ids)) @@ -575,11 +647,11 @@ def test_get_ablp_input_with_server_slice(self) -> None: ) server = dist_server.DistServer(dataset) - anchor_nodes, pos_labels, neg_labels = server.get_ablp_input( + ablp_input = server.get_ablp_input( FetchABLPInputRequest( split="train", node_type=USER, - supervision_edge_type=USER_TO_STORY, + supervision_edge_types=(USER_TO_STORY,), server_slice=ServerSlice( server_rank=0, start_numerator=1, @@ -588,6 +660,8 @@ def test_get_ablp_input_with_server_slice(self) -> None: ), ) ) + anchor_nodes = ablp_input.anchor_nodes + pos_labels, neg_labels = ablp_input.labels[USER_TO_STORY] self.assert_tensor_equality(anchor_nodes, torch.tensor([2, 3])) self.assert_tensor_equality(pos_labels, torch.tensor([[2, 3], [3, 4]]), dim=1) diff --git a/tests/unit/distributed/graph_store/messages_test.py b/tests/unit/distributed/graph_store/messages_test.py index ae8b2c0f7..cef330c51 100644 --- a/tests/unit/distributed/graph_store/messages_test.py +++ b/tests/unit/distributed/graph_store/messages_test.py @@ -1,3 +1,5 @@ +from dataclasses import FrozenInstanceError + import torch from gigl.distributed.graph_store.messages import ( @@ -5,7 +7,8 @@ FetchNodesRequest, ) from gigl.distributed.graph_store.sharding import ServerSlice -from tests.test_assets.distributed.test_dataset import USER, USER_TO_STORY +from gigl.src.common.types.graph_data import EdgeType, Relation +from tests.test_assets.distributed.test_dataset import STORY, USER, USER_TO_STORY from tests.test_assets.test_case import TestCase @@ -33,13 +36,19 @@ def test_with_server_slice(self) -> None: class TestFetchABLPInputRequest(TestCase): def test_construction(self) -> None: """Request can be constructed with required fields.""" + user_rates_story = EdgeType(USER, Relation("rates"), STORY) request = FetchABLPInputRequest( split="train", node_type=USER, - supervision_edge_type=USER_TO_STORY, + supervision_edge_types=(USER_TO_STORY, user_rates_story), ) self.assertEqual(request.split, "train") + self.assertEqual( + request.supervision_edge_types, (USER_TO_STORY, user_rates_story) + ) self.assertIsNone(request.server_slice) + with self.assertRaises(FrozenInstanceError): + request.split = "test" # ty: ignore[invalid-assignment] def test_with_server_slice(self) -> None: """Request can include a server slice.""" @@ -49,7 +58,7 @@ def test_with_server_slice(self) -> None: request = FetchABLPInputRequest( split="train", node_type=USER, - supervision_edge_type=USER_TO_STORY, + supervision_edge_types=(USER_TO_STORY,), server_slice=server_slice, ) self.assertEqual(request.server_slice, server_slice) diff --git a/tests/unit/distributed/graph_store/remote_dist_dataset_test.py b/tests/unit/distributed/graph_store/remote_dist_dataset_test.py index b1e9be293..6df121478 100644 --- a/tests/unit/distributed/graph_store/remote_dist_dataset_test.py +++ b/tests/unit/distributed/graph_store/remote_dist_dataset_test.py @@ -15,14 +15,19 @@ FetchABLPInputRequest, FetchNodesRequest, ) -from gigl.distributed.graph_store.remote_dist_dataset import RemoteDistDataset +from gigl.distributed.graph_store.remote_dist_dataset import ( + RemoteDistDataset, + _resolve_registered_supervision_edge_types, +) from gigl.distributed.graph_store.sharding import ServerSlice from gigl.env.distributed import GraphStoreInfo -from gigl.src.common.types.graph_data import EdgeType, NodeType +from gigl.src.common.types.graph_data import EdgeType, NodeType, Relation from gigl.types.graph import ( DEFAULT_HOMOGENEOUS_EDGE_TYPE, DEFAULT_HOMOGENEOUS_NODE_TYPE, FeatureInfo, + message_passing_to_negative_label, + message_passing_to_positive_label, ) from gigl.utils.sampling import ABLPInputNodes from tests.test_assets.distributed.test_dataset import ( @@ -433,11 +438,15 @@ def test_fetch_node_ids_with_splits(self, mock_async_request): torch.tensor([1, 2]), ) + @patch( + "gigl.distributed.graph_store.remote_dist_dataset.request_server", + side_effect=_mock_request_server, + ) @patch( "gigl.distributed.graph_store.remote_dist_dataset.async_request_server", side_effect=_mock_async_request_server, ) - def test_fetch_ablp_input(self, mock_async_request): + def test_fetch_ablp_input(self, mock_async_request, mock_request): """Test fetch_ablp_input with train/val/test splits.""" _create_server_with_splits() @@ -501,12 +510,16 @@ def test_fetch_ablp_input(self, mock_async_request): torch.tensor([[1]]), ) + @patch( + "gigl.distributed.graph_store.remote_dist_dataset.request_server", + side_effect=_mock_request_server, + ) @patch( "gigl.distributed.graph_store.remote_dist_dataset.async_request_server", side_effect=_mock_async_request_server, ) def test_fetch_ablp_input_respects_max_labels_per_anchor_node( - self, mock_async_request + self, mock_async_request, mock_request ): _create_server_with_splits() self.assertIsNotNone(_test_server) @@ -530,11 +543,15 @@ def test_fetch_ablp_input_respects_max_labels_per_anchor_node( torch.tensor([[2], [3], [4]]), ) + @patch( + "gigl.distributed.graph_store.remote_dist_dataset.request_server", + side_effect=_mock_request_server, + ) @patch( "gigl.distributed.graph_store.remote_dist_dataset.async_request_server", side_effect=_mock_async_request_server, ) - def test_fetch_ablp_input_with_sharding(self, mock_async_request): + def test_fetch_ablp_input_with_sharding(self, mock_async_request, mock_request): """Test fetch_ablp_input with sharding across compute nodes.""" _create_server_with_splits() @@ -648,11 +665,17 @@ def test_fetch_node_ids_auto_detects_default_node_type(self): torch.tensor([4]), ) + @patch( + "gigl.distributed.graph_store.remote_dist_dataset.request_server", + side_effect=_mock_request_server, + ) @patch( "gigl.distributed.graph_store.remote_dist_dataset.async_request_server", side_effect=_mock_async_request_server, ) - def test_fetch_ablp_input_defaults_to_homogeneous_types(self, mock_async_request): + def test_fetch_ablp_input_defaults_to_homogeneous_types( + self, mock_async_request, mock_request + ): """Test fetch_ablp_input without anchor_node_type/supervision_edge_type uses homogeneous defaults.""" self._create_labeled_homogeneous_server() cluster_info = _create_mock_graph_store_info(num_storage_nodes=1) @@ -734,6 +757,28 @@ def test_fetch_ablp_input_mismatched_params_raises(self): supervision_edge_type=DEFAULT_HOMOGENEOUS_EDGE_TYPE, ) + for supervision_edge_types in ( + [], + [USER_TO_STORY, USER_TO_STORY], + ): + with self.subTest(supervision_edge_types=supervision_edge_types): + with self.assertRaises(ValueError): + remote_dataset.fetch_ablp_input( + split="train", + anchor_node_type=USER, + supervision_edge_type=supervision_edge_types, + ) + + with self.assertRaisesRegex(ValueError, "Labeled homogeneous"): + remote_dataset.fetch_ablp_input( + split="train", + anchor_node_type=DEFAULT_HOMOGENEOUS_NODE_TYPE, + supervision_edge_type=[ + DEFAULT_HOMOGENEOUS_EDGE_TYPE, + USER_TO_STORY, + ], + ) + class TestRemoteDistDatasetSharding(RemoteDistDatasetTestBase): """Tests for fetch_node_ids and fetch_ablp_input with contiguous server assignments.""" @@ -768,17 +813,39 @@ def _mock( assert isinstance(response, torch.Tensor) response = request.server_slice.slice_tensor(response) elif isinstance(request, FetchABLPInputRequest): - if request.server_slice is not None: + if isinstance(response, ABLPInputNodes): + anchors = response.anchor_nodes + labels = response.labels + else: anchors, positive_labels, negative_labels = response - response = ( - request.server_slice.slice_tensor(anchors), - request.server_slice.slice_tensor(positive_labels), - ( - request.server_slice.slice_tensor(negative_labels) - if negative_labels is not None - else None - ), - ) + self.assertEqual(len(request.supervision_edge_types), 1) + labels = { + request.supervision_edge_types[0]: ( + positive_labels, + negative_labels, + ) + } + if request.server_slice is not None: + anchors = request.server_slice.slice_tensor(anchors) + labels = { + edge_type: ( + request.server_slice.slice_tensor(positive_labels), + ( + request.server_slice.slice_tensor(negative_labels) + if negative_labels is not None + else None + ), + ) + for edge_type, ( + positive_labels, + negative_labels, + ) in labels.items() + } + response = ABLPInputNodes( + anchor_nodes=anchors, + anchor_node_type=request.node_type, + labels=labels, + ) future.set_result(response) return future @@ -788,11 +855,17 @@ def _mock( def _mock_request_server_homogeneous( server_rank: int, func: Callable[..., Any], *args: Any, **kwargs: Any ) -> Any: - """Mock request_server that returns None for node/edge types (homogeneous).""" + """Mock homogeneous graph metadata requests.""" if func == DistServer.get_node_types: return None if func == DistServer.get_edge_types: - return None + return [ + DEFAULT_HOMOGENEOUS_EDGE_TYPE, + message_passing_to_positive_label(DEFAULT_HOMOGENEOUS_EDGE_TYPE), + message_passing_to_negative_label(DEFAULT_HOMOGENEOUS_EDGE_TYPE), + ] + if func == DistServer.get_edge_dir: + return "out" return _mock_request_server(server_rank, func, *args, **kwargs) def test_fetch_node_ids_contiguous_even_split(self) -> None: @@ -1068,7 +1141,8 @@ def test_fetch_ablp_input_contiguous_even_split(self) -> None: self.assertEqual(ablp_1.anchor_nodes.numel(), 0) pos_1, neg_1 = ablp_1.labels[DEFAULT_HOMOGENEOUS_EDGE_TYPE] self.assertEqual(pos_1.numel(), 0) - self.assertIsNone(neg_1) + assert neg_1 is not None + self.assertEqual(neg_1.numel(), 0) self.assertEqual( captured_requests, [ @@ -1077,7 +1151,7 @@ def test_fetch_ablp_input_contiguous_even_split(self) -> None: FetchABLPInputRequest( split="train", node_type=DEFAULT_HOMOGENEOUS_NODE_TYPE, - supervision_edge_type=DEFAULT_HOMOGENEOUS_EDGE_TYPE, + supervision_edge_types=(DEFAULT_HOMOGENEOUS_EDGE_TYPE,), server_slice=ServerSlice( server_rank=0, start_numerator=0, @@ -1115,7 +1189,7 @@ def test_fetch_ablp_input_contiguous_even_split(self) -> None: FetchABLPInputRequest( split="train", node_type=DEFAULT_HOMOGENEOUS_NODE_TYPE, - supervision_edge_type=DEFAULT_HOMOGENEOUS_EDGE_TYPE, + supervision_edge_types=(DEFAULT_HOMOGENEOUS_EDGE_TYPE,), server_slice=ServerSlice( server_rank=1, start_numerator=0, @@ -1127,6 +1201,193 @@ def test_fetch_ablp_input_contiguous_even_split(self) -> None: ], ) + def test_fetch_ablp_input_multiple_types_uses_one_request_per_server( + self, + ) -> None: + """Plural labels share one RPC and dense placeholders retain topology.""" + user_rates_story = EdgeType(USER, Relation("rates"), STORY) + registered_edge_types = [ + message_passing_to_positive_label(USER_TO_STORY), + message_passing_to_negative_label(USER_TO_STORY), + message_passing_to_positive_label(user_rates_story), + ] + server_data = { + 0: ABLPInputNodes( + anchor_nodes=torch.tensor([0, 1]), + anchor_node_type=USER, + labels={ + USER_TO_STORY: ( + torch.tensor([[1], [2]]), + torch.tensor([[3], [4]]), + ), + user_rates_story: (torch.tensor([[4], [3]]), None), + }, + ), + 1: ABLPInputNodes( + anchor_nodes=torch.tensor([2, 3]), + anchor_node_type=USER, + labels={ + USER_TO_STORY: ( + torch.tensor([[0], [1]]), + torch.tensor([[4], [2]]), + ), + user_rates_story: (torch.tensor([[2], [1]]), None), + }, + ), + } + captured_requests: list[tuple[int, Any]] = [] + async_mock = self._make_rank_aware_async_mock(server_data, captured_requests) + + def sync_mock( + server_rank: int, func: Callable[..., Any], *args: Any, **kwargs: Any + ) -> Any: + if func == DistServer.get_edge_dir: + return "out" + if func == DistServer.get_edge_types: + return registered_edge_types + return _mock_request_server(server_rank, func, *args, **kwargs) + + cluster_info = _create_mock_graph_store_info( + num_storage_nodes=2, num_compute_nodes=2 + ) + with _patch_remote_requests(async_mock, sync_mock): + result = RemoteDistDataset( + cluster_info=cluster_info, local_rank=0 + ).fetch_ablp_input( + split="train", + rank=0, + world_size=2, + anchor_node_type=USER, + supervision_edge_type=[USER_TO_STORY, user_rates_story], + ) + + self.assertEqual(len(captured_requests), 1) + request_rank, request = captured_requests[0] + self.assertEqual(request_rank, 0) + self.assertEqual( + request.supervision_edge_types, (USER_TO_STORY, user_rates_story) + ) + self.assertEqual(list(result[0].labels), [USER_TO_STORY, user_rates_story]) + self.assert_tensor_equality( + result[0].labels[USER_TO_STORY][0], torch.tensor([[1], [2]]) + ) + self.assertEqual(result[1].anchor_nodes.numel(), 0) + placeholder_positive, placeholder_negative = result[1].labels[USER_TO_STORY] + self.assertEqual(tuple(placeholder_positive.shape), (0, 0)) + assert placeholder_negative is not None + self.assertEqual(tuple(placeholder_negative.shape), (0, 0)) + rates_positive, rates_negative = result[1].labels[user_rates_story] + self.assertEqual(tuple(rates_positive.shape), (0, 0)) + self.assertIsNone(rates_negative) + + def test_resolve_incoming_supervision_orientations(self) -> None: + """Canonical incoming types reverse while uniform legacy types remain.""" + user_rates_story = EdgeType(USER, Relation("rates"), STORY) + registered_types = [ + message_passing_to_positive_label(STORY_TO_USER), + message_passing_to_positive_label(EdgeType(STORY, Relation("rates"), USER)), + ] + + canonical, _ = _resolve_registered_supervision_edge_types( + [USER_TO_STORY, user_rates_story], + anchor_node_type=USER, + edge_dir="in", + registered_edge_types=registered_types, + ) + self.assertEqual( + canonical, + ( + STORY_TO_USER, + EdgeType(STORY, Relation("rates"), USER), + ), + ) + legacy, _ = _resolve_registered_supervision_edge_types( + list(canonical), + anchor_node_type=USER, + edge_dir="in", + registered_edge_types=registered_types, + ) + self.assertEqual(legacy, canonical) + + with self.assertRaisesRegex(ValueError, "uniformly"): + _resolve_registered_supervision_edge_types( + [USER_TO_STORY, EdgeType(STORY, Relation("rates"), USER)], + anchor_node_type=USER, + edge_dir="in", + registered_edge_types=registered_types, + ) + + def test_fetch_ablp_input_multiple_types_slices_every_label_row(self) -> None: + """Fractional server slicing applies the same rows to every objective.""" + user_rates_story = EdgeType(USER, Relation("rates"), STORY) + server_data = { + 0: ABLPInputNodes( + anchor_nodes=torch.tensor([0, 1, 2, 3]), + anchor_node_type=USER, + labels={ + USER_TO_STORY: ( + torch.tensor([[10], [11], [12], [13]]), + torch.tensor([[20], [21], [22], [23]]), + ), + user_rates_story: ( + torch.tensor([[30], [31], [32], [33]]), + None, + ), + }, + ) + } + captured_requests: list[tuple[int, Any]] = [] + async_mock = self._make_rank_aware_async_mock(server_data, captured_requests) + registered_edge_types = [ + message_passing_to_positive_label(USER_TO_STORY), + message_passing_to_negative_label(USER_TO_STORY), + message_passing_to_positive_label(user_rates_story), + ] + + def sync_mock( + server_rank: int, func: Callable[..., Any], *args: Any, **kwargs: Any + ) -> Any: + if func == DistServer.get_edge_dir: + return "out" + if func == DistServer.get_edge_types: + return registered_edge_types + return _mock_request_server(server_rank, func, *args, **kwargs) + + cluster_info = _create_mock_graph_store_info( + num_storage_nodes=1, num_compute_nodes=2 + ) + with _patch_remote_requests(async_mock, sync_mock): + result = RemoteDistDataset( + cluster_info=cluster_info, local_rank=0 + ).fetch_ablp_input( + split="train", + rank=1, + world_size=2, + anchor_node_type=USER, + supervision_edge_type=[USER_TO_STORY, user_rates_story], + ) + + self.assert_tensor_equality(result[0].anchor_nodes, torch.tensor([2, 3])) + self.assert_tensor_equality( + result[0].labels[USER_TO_STORY][0], torch.tensor([[12], [13]]) + ) + negative_labels = result[0].labels[USER_TO_STORY][1] + assert negative_labels is not None + self.assert_tensor_equality(negative_labels, torch.tensor([[22], [23]])) + self.assert_tensor_equality( + result[0].labels[user_rates_story][0], torch.tensor([[32], [33]]) + ) + self.assertEqual(len(captured_requests), 1) + self.assertEqual( + captured_requests[0][1].server_slice, + ServerSlice( + server_rank=0, + start_numerator=1, + end_numerator=2, + denominator=2, + ), + ) + def test_fetch_ablp_input_contiguous_fractional_split(self) -> None: """ABLP CONTIGUOUS with 3 storage nodes and 2 compute nodes: server 1 fractionally split.""" server_data: dict[ @@ -1189,7 +1450,7 @@ def test_fetch_ablp_input_contiguous_fractional_split(self) -> None: FetchABLPInputRequest( split="train", node_type=DEFAULT_HOMOGENEOUS_NODE_TYPE, - supervision_edge_type=DEFAULT_HOMOGENEOUS_EDGE_TYPE, + supervision_edge_types=(DEFAULT_HOMOGENEOUS_EDGE_TYPE,), server_slice=ServerSlice( server_rank=0, start_numerator=0, @@ -1203,7 +1464,7 @@ def test_fetch_ablp_input_contiguous_fractional_split(self) -> None: FetchABLPInputRequest( split="train", node_type=DEFAULT_HOMOGENEOUS_NODE_TYPE, - supervision_edge_type=DEFAULT_HOMOGENEOUS_EDGE_TYPE, + supervision_edge_types=(DEFAULT_HOMOGENEOUS_EDGE_TYPE,), server_slice=ServerSlice( server_rank=1, start_numerator=0, @@ -1250,7 +1511,7 @@ def test_fetch_ablp_input_contiguous_fractional_split(self) -> None: FetchABLPInputRequest( split="train", node_type=DEFAULT_HOMOGENEOUS_NODE_TYPE, - supervision_edge_type=DEFAULT_HOMOGENEOUS_EDGE_TYPE, + supervision_edge_types=(DEFAULT_HOMOGENEOUS_EDGE_TYPE,), server_slice=ServerSlice( server_rank=1, start_numerator=1, @@ -1264,7 +1525,7 @@ def test_fetch_ablp_input_contiguous_fractional_split(self) -> None: FetchABLPInputRequest( split="train", node_type=DEFAULT_HOMOGENEOUS_NODE_TYPE, - supervision_edge_type=DEFAULT_HOMOGENEOUS_EDGE_TYPE, + supervision_edge_types=(DEFAULT_HOMOGENEOUS_EDGE_TYPE,), server_slice=ServerSlice( server_rank=2, start_numerator=0, diff --git a/tests/unit/distributed/label_remap_cuda_device_test.py b/tests/unit/distributed/label_remap_cuda_device_test.py new file mode 100644 index 000000000..53ac1ff07 --- /dev/null +++ b/tests/unit/distributed/label_remap_cuda_device_test.py @@ -0,0 +1,60 @@ +"""CUDA device-placement regression test for ABLP label remapping.""" + +import unittest + +import torch + +from gigl.distributed.utils.ablp import remap_labels_to_local_edge_indices +from gigl.src.common.types.graph_data import EdgeType, NodeType, Relation +from gigl.types.graph import message_passing_to_positive_label +from tests.test_assets.test_case import TestCase + +_USER = NodeType("user") +_STORY = NodeType("story") +_USER_TO_STORY = EdgeType(_USER, Relation("to"), _STORY) + + +def _inputs( + device: torch.device, +) -> tuple[dict[NodeType, torch.Tensor], dict[EdgeType, torch.Tensor]]: + """Build an unsorted node map and padded labels on one device.""" + node_map = {_STORY: torch.tensor([15, 10, 16, 11, 12], device=device)} + positive_labels = { + message_passing_to_positive_label(_USER_TO_STORY): torch.tensor( + [[15, 15], [16, -1], [-1, -1]], + dtype=torch.long, + device=device, + ) + } + return node_map, positive_labels + + +@unittest.skipUnless(torch.cuda.is_available(), "requires a CUDA device") +class LabelRemapCudaDeviceTest(TestCase): + def test_cuda_output_matches_cpu(self) -> None: + cpu_node_map, cpu_positive_labels = _inputs(torch.device("cpu")) + expected_positive, _ = remap_labels_to_local_edge_indices( + local_id_to_global_id_by_node_type=cpu_node_map, + positive_labels_by_edge_type=cpu_positive_labels, + negative_labels_by_edge_type={}, + ) + + cuda_node_map, cuda_positive_labels = _inputs(torch.device("cuda")) + actual_positive, _ = remap_labels_to_local_edge_indices( + local_id_to_global_id_by_node_type=cuda_node_map, + positive_labels_by_edge_type=cuda_positive_labels, + negative_labels_by_edge_type={}, + ) + + self.assertEqual(set(actual_positive.keys()), set(expected_positive.keys())) + for edge_type, expected_edge_index in expected_positive.items(): + actual_edge_index = actual_positive[edge_type] + self.assertEqual(actual_edge_index.device.type, "cuda") + torch.testing.assert_close( + actual_edge_index.cpu(), + expected_edge_index, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/distributed/utils/ablp_test.py b/tests/unit/distributed/utils/ablp_test.py new file mode 100644 index 000000000..a65726d65 --- /dev/null +++ b/tests/unit/distributed/utils/ablp_test.py @@ -0,0 +1,285 @@ +"""Tests for vectorized ABLP label remapping.""" + +import unittest + +import torch +from parameterized import param, parameterized +from torch_geometric.typing import EdgeType as PyGEdgeType + +from gigl.distributed.utils.ablp import ( + label_edge_index_to_dict, + remap_labels_to_local_edge_indices, +) +from gigl.src.common.types.graph_data import EdgeType, NodeType, Relation +from gigl.types.graph import ( + DEFAULT_HOMOGENEOUS_EDGE_TYPE, + DEFAULT_HOMOGENEOUS_NODE_TYPE, + message_passing_to_negative_label, + message_passing_to_positive_label, +) +from tests.test_assets.test_case import TestCase + +_USER = NodeType("user") +_STORY = NodeType("story") +_USER_TO_STORY = EdgeType(_USER, Relation("to"), _STORY) +_A = NodeType("a") +_B = NodeType("b") +_C = NodeType("c") +_A_TO_B = EdgeType(_A, Relation("to"), _B) +_A_TO_C = EdgeType(_A, Relation("to"), _C) + + +def _positive_label_edge_type(edge_type: EdgeType) -> EdgeType: + return message_passing_to_positive_label(edge_type) + + +def _negative_label_edge_type(edge_type: EdgeType) -> EdgeType: + return message_passing_to_negative_label(edge_type) + + +def _assert_label_sets_equal( + actual: dict[PyGEdgeType, torch.Tensor], + expected: dict[PyGEdgeType, dict[int, list[int]]], +) -> None: + assert set(actual.keys()) == set(expected.keys()) + for edge_type, expected_by_anchor in expected.items(): + actual_by_anchor = label_edge_index_to_dict( + label_edge_index=actual[edge_type], + num_anchors=len(expected_by_anchor), + ) + assert set(actual_by_anchor.keys()) == set(expected_by_anchor.keys()) + for anchor, expected_labels in expected_by_anchor.items(): + actual_labels = actual_by_anchor[anchor] + assert actual_labels.dtype == torch.long + assert sorted(actual_labels.tolist()) == sorted(expected_labels) + + +class LabelEdgeIndexToDictTest(TestCase): + def test_preserves_empty_and_multi_label_anchors(self) -> None: + label_edge_index = torch.tensor([[0, 2, 2], [5, 7, 8]]) + + labels_by_anchor = label_edge_index_to_dict( + label_edge_index=label_edge_index, + num_anchors=3, + ) + + self.assertEqual(set(labels_by_anchor.keys()), {0, 1, 2}) + torch.testing.assert_close( + labels_by_anchor[0], torch.tensor([5], dtype=torch.long) + ) + torch.testing.assert_close( + labels_by_anchor[1], torch.empty(0, dtype=torch.long) + ) + torch.testing.assert_close( + labels_by_anchor[2], torch.tensor([7, 8], dtype=torch.long) + ) + + def test_all_anchors_empty(self) -> None: + labels_by_anchor = label_edge_index_to_dict( + label_edge_index=torch.empty((2, 0), dtype=torch.long), + num_anchors=2, + ) + + self.assertEqual(set(labels_by_anchor.keys()), {0, 1}) + torch.testing.assert_close( + labels_by_anchor[0], torch.empty(0, dtype=torch.long) + ) + torch.testing.assert_close( + labels_by_anchor[1], torch.empty(0, dtype=torch.long) + ) + + +class RemapLabelsToLocalEdgeIndicesTest(TestCase): + @parameterized.expand( + [ + param( + "sorted_present_empty_and_padded", + local_id_to_global_id_by_node_type={ + _STORY: torch.tensor([10, 11, 12, 13, 14, 15, 16, 17]) + }, + positive_labels={ + _positive_label_edge_type(_USER_TO_STORY): torch.tensor( + [[15, -1], [15, 16], [-1, -1], [99, -1]] + ) + }, + negative_labels={}, + expected_positive={_USER_TO_STORY: {0: [5], 1: [5, 6], 2: [], 3: []}}, + expected_negative={}, + ), + param( + "duplicate_labels_preserve_multiplicity", + local_id_to_global_id_by_node_type={ + _STORY: torch.tensor([10, 11, 12, 13, 14, 15]) + }, + positive_labels={ + _positive_label_edge_type(_USER_TO_STORY): torch.tensor( + [[15, 15], [11, 11]] + ) + }, + negative_labels={}, + expected_positive={_USER_TO_STORY: {0: [5, 5], 1: [1, 1]}}, + expected_negative={}, + ), + param( + "default_homogeneous_keying", + local_id_to_global_id_by_node_type={ + DEFAULT_HOMOGENEOUS_NODE_TYPE: torch.tensor([20, 10, 30, 11, 15]) + }, + positive_labels={ + message_passing_to_positive_label( + DEFAULT_HOMOGENEOUS_EDGE_TYPE + ): torch.tensor([[30, 10], [-1, -1]]) + }, + negative_labels={}, + expected_positive={DEFAULT_HOMOGENEOUS_EDGE_TYPE: {0: [2, 1], 1: []}}, + expected_negative={}, + ), + param( + "positive_and_negative_labels", + local_id_to_global_id_by_node_type={ + _STORY: torch.tensor([10, 11, 12, 13, 14, 15, 16, 17]) + }, + positive_labels={ + _positive_label_edge_type(_USER_TO_STORY): torch.tensor( + [[15], [16]] + ) + }, + negative_labels={ + _negative_label_edge_type(_USER_TO_STORY): torch.tensor( + [[13, 16], [17, -1]] + ) + }, + expected_positive={_USER_TO_STORY: {0: [5], 1: [6]}}, + expected_negative={_USER_TO_STORY: {0: [3, 6], 1: [7]}}, + ), + param( + "multiple_supervision_edge_types", + local_id_to_global_id_by_node_type={ + _B: torch.tensor([11, 12, 13, 14, 15, 16]), + _C: torch.tensor([20, 21, 22, 23, 24, 25]), + }, + positive_labels={ + _positive_label_edge_type(_A_TO_B): torch.tensor([[13, 14]]), + _positive_label_edge_type(_A_TO_C): torch.tensor([[22, 23]]), + }, + negative_labels={ + _negative_label_edge_type(_A_TO_B): torch.tensor([[15, 16]]), + _negative_label_edge_type(_A_TO_C): torch.tensor([[24, 25]]), + }, + expected_positive={ + _A_TO_B: {0: [2, 3]}, + _A_TO_C: {0: [2, 3]}, + }, + expected_negative={ + _A_TO_B: {0: [4, 5]}, + _A_TO_C: {0: [4, 5]}, + }, + ), + param( + "all_anchors_empty", + local_id_to_global_id_by_node_type={_STORY: torch.tensor([10, 11, 12])}, + positive_labels={ + _positive_label_edge_type(_USER_TO_STORY): torch.tensor( + [[-1, -1], [99, 98]] + ) + }, + negative_labels={}, + expected_positive={_USER_TO_STORY: {0: [], 1: []}}, + expected_negative={}, + ), + ] + ) + def test_matches_constructed_labels( + self, + _, + local_id_to_global_id_by_node_type: dict[NodeType, torch.Tensor], + positive_labels: dict[EdgeType, torch.Tensor], + negative_labels: dict[EdgeType, torch.Tensor], + expected_positive: dict[PyGEdgeType, dict[int, list[int]]], + expected_negative: dict[PyGEdgeType, dict[int, list[int]]], + ) -> None: + positive_edge_indices, negative_edge_indices = ( + remap_labels_to_local_edge_indices( + local_id_to_global_id_by_node_type=local_id_to_global_id_by_node_type, + positive_labels_by_edge_type=positive_labels, + negative_labels_by_edge_type=negative_labels, + ) + ) + + _assert_label_sets_equal(positive_edge_indices, expected_positive) + _assert_label_sets_equal(negative_edge_indices, expected_negative) + + def test_pins_exact_edge_index(self) -> None: + positive_edge_indices, _ = remap_labels_to_local_edge_indices( + local_id_to_global_id_by_node_type={ + _STORY: torch.tensor([10, 11, 12, 13, 14, 15, 16, 17]) + }, + positive_labels_by_edge_type={ + _positive_label_edge_type(_USER_TO_STORY): torch.tensor( + [[15, -1], [15, 16], [-1, -1]] + ) + }, + negative_labels_by_edge_type={}, + ) + + torch.testing.assert_close( + positive_edge_indices[_USER_TO_STORY], + torch.tensor([[0, 1, 1], [5, 5, 6]], dtype=torch.long), + ) + + def test_unsorted_node_map_returns_local_indices(self) -> None: + positive_edge_indices, _ = remap_labels_to_local_edge_indices( + local_id_to_global_id_by_node_type={_STORY: torch.tensor([15, 10, 16, 11])}, + positive_labels_by_edge_type={ + _positive_label_edge_type(_USER_TO_STORY): torch.tensor([[16, 15]]) + }, + negative_labels_by_edge_type={}, + ) + + labels_by_anchor = label_edge_index_to_dict( + positive_edge_indices[_USER_TO_STORY], num_anchors=1 + ) + self.assertEqual(sorted(labels_by_anchor[0].tolist()), [0, 2]) + + def test_zero_anchor_tensor_yields_no_edge_type_key(self) -> None: + positive_edge_indices, negative_edge_indices = ( + remap_labels_to_local_edge_indices( + local_id_to_global_id_by_node_type={_STORY: torch.tensor([10, 11, 12])}, + positive_labels_by_edge_type={ + _positive_label_edge_type(_USER_TO_STORY): torch.empty((0, 0)) + }, + negative_labels_by_edge_type={}, + ) + ) + + self.assertEqual(positive_edge_indices, {}) + self.assertEqual(negative_edge_indices, {}) + + def test_output_follows_input_device_and_dtype(self) -> None: + positive_edge_indices, _ = remap_labels_to_local_edge_indices( + local_id_to_global_id_by_node_type={ + _STORY: torch.tensor([10, 11, 12, 13, 14, 15]) + }, + positive_labels_by_edge_type={ + _positive_label_edge_type(_USER_TO_STORY): torch.tensor([[15], [11]]) + }, + negative_labels_by_edge_type={}, + ) + + label_edge_index = positive_edge_indices[_USER_TO_STORY] + self.assertEqual(label_edge_index.device.type, "cpu") + self.assertEqual(label_edge_index.dtype, torch.long) + + def test_duplicate_node_map_raises(self) -> None: + with self.assertRaises(ValueError): + remap_labels_to_local_edge_indices( + local_id_to_global_id_by_node_type={_STORY: torch.tensor([10, 10, 11])}, + positive_labels_by_edge_type={ + _positive_label_edge_type(_USER_TO_STORY): torch.tensor([[10, 11]]) + }, + negative_labels_by_edge_type={}, + ) + + +if __name__ == "__main__": + unittest.main()