Implementing NodeSelectors and Tolerations with TorchX on Kubernetes(EKS)

username: vara_bonthu
cooked: <p>I’m currently working with TorchX in conjunction with Volcano scheduling for my training jobs on an Amazon EKS cluster. I’ve also integrated Karpenter autoscaler for effective node scaling. Additionally, I’m using managed node groups with labeled nodes that have specific taints applied.</p>
<p>Our internal data and machine learning teams have the requirement to specify NodeSelectors and Tolerations to target jobs on particular nodes or managed node groups. While referring to the documentation provided here: <a href="https://pytorch.org/torchx/main/specs.html" rel="noopener nofollow ugc">TorchX Specifications</a>, I observed that <em>capabilities={“<a href="http://node.kubernetes.io/instance-type" rel="noopener nofollow ugc">node.kubernetes.io/instance-type</a>”: “”}</em> are used as NodeSelectors when the job is created through Volcano. However, this approach doesn’t seem to allow for sending a list of labels, which our use case demands.</p>
<p>Furthermore, I’m also interested in incorporating tolerations into these jobs to ensure proper scheduling and execution in our environment. If any of you have experience in implementing NodeSelectors and Tolerations in TorchX within an Amazon EKS setup, I would highly appreciate your insights and advice.</p>
<p>If there’s no previous experience with this scenario, I’m considering raising a feature request to address these needs. Your guidance and input would be greatly valued.</p>

username: lucifer177
cooked: <aside class="quote no-group" data-username="vara_bonthu" data-post="1" data-topic="186437">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/vara_bonthu/48/63200_2.png" class="avatar"> vara_bonthu:</div>
<blockquote>
<p>Our internal data and machine learning teams have the requirement to specify NodeSelectors and Tolerations to target jobs on particular nodes or managed node groups. While referring to the documentation provided here: <a href="https://pytorch.org/torchx/main/specs.html" rel="noopener nofollow ugc">TorchX Specifications </a>, I observed that <em>capabilities={“<a href="http://node.kubernetes.io/instance-type" rel="noopener nofollow ugc">node.kubernetes.io/instance-type </a>”: “”}</em> are used as NodeSelectors when the job is created through Volcano. However, this approach doesn’t seem to allow for sending a list of labels, which our use case demands.</p>
</blockquote>
</aside>
<p>Did you find a solution for this?<br>
Additionally, do you use any autoscaling such as Karpenter or Cluster AutoScaler?</p>

username: steph_dev
cooked: <p><code>capabilities</code> config works but its not enough for our use case…</p>
<p>hope they implement “tolerations” config to schedule nodes on tainted nodes.</p>

username: python
cooked: <p>I’ve run into something similar with TorchX on EKS. As you mentioned, the <code>capabilities</code> field maps to NodeSelectors, but it’s pretty limited; it doesn’t handle multiple labels or tolerations natively. What worked for us was customizing the scheduler or extending the TorchX component to add the extra Kubernetes fields manually before submission.</p>
<p>If you’re using a custom scheduler or running TorchX through a wrapper script, you could patch the generated pod spec using the Volcano Job API or a mutating webhook to inject the NodeSelectors and Tolerations. Not ideal, but it gives you full control.</p>
<p>I’d support a feature request for better native support for these Kubernetes features in TorchX would help a lot in more complex EKS setups like yours.</p>