Codex/graph store multi supervision edge types

[kmontemayor2-sc]

Scope of work done

Where is the documentation for this feature?: N/A

Did you add automated tests or write a test plan?

Updated Changelog.md? NO

Ready for code review?: NO

[kmontemayor2-sc]

/all_test

[github-actions]

GiGL Automation

@ 20:06:49UTC : 🔄 Integration Test started.

@ 21:40:29UTC : ✅ Workflow completed successfully.

[github-actions]

GiGL Automation

@ 20:06:50UTC : 🔄 E2E Test started.

[github-actions]

GiGL Automation

@ 20:06:53UTC : 🔄 Python Unit Test started.

@ 21:31:02UTC : ❌ Workflow failed.
Please check the logs for more details.

[github-actions]

GiGL Automation

@ 20:06:55UTC : 🔄 Scala Unit Test started.

@ 20:14:57UTC : ✅ Workflow completed successfully.

[github-actions]

GiGL Automation

@ 20:06:55UTC : 🔄 Lint Test started.

@ 20:15:23UTC : ✅ Workflow completed successfully.

[github-actions]

GiGL Automation

@ 20:06:56UTC : 🔄 C++ Unit Test started.

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

[semgrep-code-snapchat]

Semgrep identified an issue in your code:

Pickle-based object serialization in all_gather_object() can execute attacker-controlled code on remote ranks if validation errors are compromised.

More details about this

The torch.distributed.all_gather_object() call uses pickle serialization under the hood to share the validation_error string across distributed ranks. If an attacker can control the content of validation_error, they can craft a malicious pickled object that executes arbitrary code when deserialized on other ranks.

Exploit scenario:

1. An attacker triggers a ValueError in _convert_graph_store_ablp_inputs() by providing malicious input (e.g., crafted node data, edge types, or edge direction that causes a specific error condition).
2. The error message is converted to a string and stored in validation_error.
3. When all_gather_object() serializes this string using pickle and sends it to other ranks, if the attacker has injected a specially-crafted pickle payload instead of a plain string, the deserialization on remote ranks will execute the attacker's code.
4. The attacker gains code execution on all compute ranks in the distributed training job.

Even though str(error) is called here, if an exception object with a malicious __str__ method reaches this point, or if the exception's string representation is later pickled as part of a larger object, pickle exploitation becomes possible.

To resolve this comment:

✨ Commit fix suggestion

1. Replace torch.distributed.all_gather_object(validation_errors, validation_error) with tensor-based collectives so PyTorch does not use pickle under the hood.

2. Gather the error state as an integer tensor instead of a Python string.
   For example, create a local flag like has_error = torch.tensor([1 if validation_error is not None else 0], device=...) and gather it with torch.distributed.all_gather(...) into one tensor per rank.

3. Change the validation check to use the gathered flags.
   For example, raise when any gathered value is 1 instead of checking a list of Python objects.

4. Keep the exception message generic, or send only a numeric error code if you need rank-specific detail.
   For example, use error_code = torch.tensor([0 if validation_error is None else 1], device=...) and map codes to messages locally after all_gather(...).

5. If you must preserve the text of the validation error, encode it into bytes and gather it with tensors instead of object collectives.
   For example, convert the message to utf-8, store its length in an integer tensor, store the bytes in a fixed-size torch.uint8 tensor, and call torch.distributed.all_gather(...) on those tensors. This avoids deserializing arbitrary Python objects from other ranks.

6. Update the final ValueError to use the safe gathered result.
   For example, raise ValueError("Graph Store ABLP input validation failed across compute ranks.") or reconstruct messages only from tensor data gathered in the previous step.

💬 Ignore this finding

Reply with Semgrep commands to ignore this finding.

• /fp <comment> for false positive
• /ar <comment> for acceptable risk
• /other <comment> for all other reasons

Alternatively, triage in Semgrep AppSec Platform to ignore the finding created by pickles-in-pytorch-distributed.

_{You can view more details about this finding in the Semgrep AppSec Platform.}