Kmonte/shared backend decomp 4

[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

@ 21:22:55UTC : 🔄 Integration Test started.

@ 22:40:25UTC : ✅ Workflow completed successfully.

[github-actions]

GiGL Automation

@ 21:22:57UTC : 🔄 Python Unit Test started.

@ 22:32:01UTC : ✅ Workflow completed successfully.

[github-actions]

GiGL Automation

@ 21:22:59UTC : 🔄 Scala Unit Test started.

@ 21:32:06UTC : ✅ Workflow completed successfully.

[github-actions]

GiGL Automation

@ 21:22:59UTC : 🔄 E2E Test started.

@ 22:47:31UTC : ✅ Workflow completed successfully.

[github-actions]

GiGL Automation

@ 21:23:00UTC : 🔄 Lint Test started.

@ 21:31:11UTC : ✅ Workflow completed successfully.

[semgrep-code-snapchat]

Semgrep identified an issue in your code:

The torch.distributed.all_gather_object() call deserializes untrusted data from remote ranks using pickle, allowing arbitrary code execution if an attacker controls any rank in the distributed system.

More details about this

The torch.distributed.all_gather_object() call is using pickle deserialization under the hood to share producer_id_list across all ranks in the distributed system. This means untrusted data from other processes gets automatically deserialized and executed.

An attacker with access to any rank in the distributed training job could craft a malicious pickled object and send it during the all-gather phase. When producer_id_list (or other ranks' data) gets deserialized on your rank, the attacker's code runs immediately with the same privileges as your training process.

For example:

1. Attacker gains write access to rank 1's memory or intercepts its state
2. They insert a pickled Python object that executes shell commands when unpickled (e.g., os.system('steal_data.sh'))
3. Your rank calls torch.distributed.all_gather_object(all_producer_ids, producer_id_list)
4. PyTorch's pickle unpickles all ranks' data, triggering the attacker's code during deserialization on your process
5. The shell commands run with your process's credentials, potentially exfiltrating model weights or training data

To resolve this comment:

✨ Commit Assistant Fix Suggestion

1. Avoid using torch.distributed.all_gather_object, as this relies on Python pickling, which may lead to arbitrary code execution if untrusted data is ever deserialized.
2. Replace the use of all_gather_object with a tensor-based collective, such as torch.distributed.all_gather, by converting your data to a tensor (for example, use torch.tensor(producer_id_list)).
3. Predefine the size and type of the tensor holding the gathered data, for example: all_producer_ids = torch.empty(runtime.world_size * num_producers, dtype=torch.long), where num_producers is the expected length of producer_id_list for each rank.
4. Call torch.distributed.all_gather([all_producer_ids], producer_id_list_tensor), where producer_id_list_tensor = torch.tensor(producer_id_list, dtype=torch.long) for each rank.
5. After gathering, reconstruct the original data structure as needed from all_producer_ids. For example, split the combined tensor into per-rank lists using slicing.
6. If the number of producers may vary, agree on a fixed length and pad with a sentinel value such as -1 so tensors can be safely communicated.

This change ensures only primitive tensor data is shared between ranks, eliminating pickle-related risks.

💬 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.}

[semgrep-code-snapchat]

Semgrep identified an issue in your code:

Using torch.distributed.all_gather_object() to share my_worker_key creates an arbitrary code execution risk, since pickle deserialization can execute attacker-controlled code if a compromised worker sends malicious data.

More details about this

The torch.distributed.all_gather_object() call on this line uses pickle to serialize and deserialize the my_worker_key string across distributed processes. An attacker who can control the data sent from any worker process could craft a malicious pickle payload that executes arbitrary code when deserialized.

Exploit scenario:

1. An attacker compromises or spoofs one of the worker processes in the distributed training cluster
2. They set my_worker_key to a malicious pickle-serialized object instead of a normal string
3. When this line executes, PyTorch deserializes the pickle object on all receiving ranks
4. The malicious pickle payload executes arbitrary code on those processes with the same privileges as the training job, potentially stealing model weights, injecting backdoors, or exfiltrating data

To resolve this comment:

✨ Commit Assistant Fix Suggestion

1. Avoid using torch.distributed.all_gather_object, as it uses pickle internally and can allow arbitrary code execution if untrusted data is deserialized.
2. If exchanging strings across ranks, switch to using torch.distributed.all_gather, which is safe for tensors. You can do this by encoding your strings to byte tensors before gathering and decoding after.
3. Replace the vulnerable line with logic similar to:
   • Convert the local string to bytes: my_worker_key_bytes = my_worker_key.encode('utf-8')
   • Find the maximum length of all keys to ensure tensors are the same size across ranks. This usually requires an all_reduce to get the max. For example:
     key_len_tensor = torch.tensor([len(my_worker_key_bytes)], device='cpu')
max_len_tensor = key_len_tensor.clone()
torch.distributed.all_reduce(max_len_tensor, op=torch.distributed.ReduceOp.MAX)
   • Pad your byte string to max_len_tensor.item(): padded_bytes = my_worker_key_bytes.ljust(max_len_tensor.item(), b'\x00')
   • Create a tensor: my_worker_key_tensor = torch.ByteTensor(list(padded_bytes))
   • Prepare a gather tensor: all_worker_key_tensors = [torch.empty_like(my_worker_key_tensor) for _ in range(runtime.world_size)]
   • Call torch.distributed.all_gather(all_worker_key_tensors, my_worker_key_tensor)
   • After gathering, decode each tensor: [bytes(t.tolist()).rstrip(b'\x00').decode('utf-8') for t in all_worker_key_tensors]
4. Replace all uses of all_worker_keys with the decoded string list.

Alternatively, if all ranks already know or can deterministically construct the set of worker keys, you can avoid broadcasting entirely by constructing the list locally.

Using tensors for communication prevents vulnerabilities from deserialization attacks, as tensor operations do not use pickle.

💬 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.}