Running Ray Job as a Flyte Task

1

Hello,

I am triggering typic Ray Batch inference flow as a flyte task. In the flyte task I mention a container image that has pytorch version 2.6 lib dependency.

Ray job invokes a Ray Actor on each node in Ray Cluster that does the pytorch inference. Code for PyTorch Ray Actor can directly do “import torch” because torch is in the container image.

What I want to do is, I receive dynamic input before this flyte task that tell me Torch version,lets say 2.5. I want my PyTorch Ray Actor to do the inference on version 2.5 and not the 2.6 baked in the container image of flyte task.

How do I do that?

2

Hi Monil!

You can try to specify a different PyTorch version for your Ray Actor by using Ray’s runtime_env feature with the pip field, which allows you to override or add Python packages at runtime.

For example, you can set runtime_env={"pip": ["torch==2.5.0"]} when creating the actor, and Ray will install the specified version in an isolated environment for that actor, I think this should hypothetically work for your use case :slight_smile:

You can read more about it here: Environment Dependencies — Ray 2.49.1