[Bug]: Pytorch RunInference PredictionResult is a Dict

[yeandy]

What happened?

The return value of the forward call for many models is a dictionary with the predictions along with more metadata. i.e. Dict[str, Tensor]. However, RunInference currently expects outputs to be an Iterable[Any]. i.e. Iterable[Tensor] or Iterable[Dict[str, Tensor]]. So when RunInference zips the inputs with the predictions, the predictions iterate over the dictionary keys instead of the batch elements. So we end up with only the key name and actual prediction Tensors are discarded. This results in an unusable PredictionResult.

The PredictionResult challenge is a known quirk due to the current design/implementation of RunInference. We have a sample that does something similar by creating a wrapper class. We should support both tensors and dictionary of tensors. Currently, it only works for the former, and that is baked into the logic. We should figure out how to natively support results that are a dict type.

Issue Priority

Priority: 2

Issue Component

Component: sdk-py-core

[yeandy]

.take-issue

[yeandy]

Subtask of #22117

[damccorm]

#22979 got created as an accidental duplicate. I'm closing that one, but copying over the description since I think it has some extra additional context:

Originally documented here in our RunInference Troubleshooting section:

For some models, the PredictionResults output might not include the correct predictions in the inferences field. This issue occurs when you use a model whose inferences return a dictionary that maps keys to predictions and other metadata. An example return type is Dict[str, Tensor].

The RunInference API currently expects outputs to be an Iterable[Any]. Example return types are Iterable[Tensor] or Iterable[Dict[str, Tensor]]. When RunInference zips the inputs with the predictions, the predictions iterate over the dictionary keys instead of the batch elements. The result is that the key name is preserved but the prediction tensors are discarded.

To work with the current RunInference implementation, you can create a wrapper class that overrides the model(input) call. In PyTorch, for example, your wrapper would override the forward() function and return an output with the appropriate format of List[Dict[str, torch.Tensor]].

[yeandy]

Thanks @damccorm! We should also remove (or at least clarify) the section on this issue from our Troubleshooting guide. I had intended to make this comment in the PR yesterday, but forgot, and just remembered. 😄

[damccorm]

That's a good callout - I'll reopen here. Since this fix won't be released for ~2 months (it missed the 2.42 cut), we should probably hold off on doing that until then.

[damccorm]

#23152 will close this once its merged (after the 2.43 release)