Add trainium, inferentia, and efa parameters to @kubernetes decorator

[emattia]

PR Type

• Bug fix
• New feature
• Core Runtime change
• Docs / tooling
• Refactoring

Summary

Mirror @batch's AWS-accelerator surface on @kubernetes:

• @kubernetes(trainium=N) requests N AWS Trainium / Inferentia Neuron
  devices (aws.amazon.com/neuron k8s resource).
• @kubernetes(inferentia=N) is an alias for trainium, mirroring
  @batch(inferentia=N) for API consistency.
• @kubernetes(efa=N) requests N AWS Elastic Fabric Adapter network
  interfaces (vpc.amazonaws.com/efa k8s resource).

Plumbed through kubernetes_job, kubernetes_jobsets, kubernetes_cli,
and the argo / airflow runtimes consistently with how the existing
gpu parameter is handled.

Issue

No tracking issue. Supersedes the original PR scope of just trainium.
Brings the @kubernetes path to parity with @batch for AWS Neuron
and EFA workloads, unblocking customers who run their own EKS clusters
and want first-class Neuron/EFA support without writing raw pod specs.

Reproduction

Runtime: kubernetes (EKS with AWS Neuron and EFA device plugins installed; nodes labeled with the relevant accelerator).

Commands to run:

from metaflow import FlowSpec, step, kubernetes, environment

NEURON_IMG = "public.ecr.aws/neuron/pytorch-training-neuronx:2.9.0-neuronx-py312-sdk2.29.1-ubuntu24.04"

class NeuronEfaSmoke(FlowSpec):

    @kubernetes(trainium=1, image=NEURON_IMG)
    @step
    def neuron_only(self):
        import subprocess
        print(subprocess.check_output(["neuron-ls"]).decode())
        self.next(self.gpu_efa)

    # Equivalent - inferentia is an alias for trainium
    @kubernetes(inferentia=1, image=NEURON_IMG)
    @step
    def inferentia_alias(self):
        ...

    @environment(vars={"FI_PROVIDER": "efa"})
    @kubernetes(gpu=8, efa=32, image="<aws-dlc-pytorch-cuda>")
    @step
    def gpu_with_efa(self):
        import torch.distributed as dist
        dist.init_process_group(backend="nccl")
        # NCCL debug log will show "Selected provider is efa"

Where evidence shows up: task pod spec (kubectl describe pod) and
NCCL debug log inside the running container.

Before (master)

TypeError: kubernetes() got an unexpected keyword argument 'trainium'

(also for inferentia, efa)

After (this PR)

$ kubectl describe pod ws-...
...
Limits:
  aws.amazon.com/neuron:    1
  vpc.amazonaws.com/efa:    32

Root Cause

Not a bug fix — net-new feature. The underlying Kubernetes resources
(aws.amazon.com/neuron, vpc.amazonaws.com/efa) are advertised by the
respective AWS device plugins; @kubernetes had no decorator-level
surface to request them. @batch already exposed trainium,
inferentia, and efa. This PR brings @kubernetes to parity.

Why This Fix Is Correct

• Mirrors @batch's API surface exactly. inferentia collapses into
  trainium at step_init and is popped before any runtime translation
  — same shape as batch_decorator.py:175-211, only with trainium as
  canonical (since on K8s the underlying resource name is
  aws.amazon.com/neuron and we surface what users running on Trainium
  hardware naturally type first).
• Doesn't disturb the existing GPU path. gpu and trainium are
  enforced as mutually exclusive (matching @batch's convention).
• Argo/Airflow runtimes already had the trainium plumbing pattern from
  earlier in this branch; efa follows the same pattern.

Failure Modes Considered

1. Backward compat: flows using only gpu / gpu_vendor are
   unaffected — new attributes default to None and resource-limit
   emission is gated on non-None values.
2. Mutual exclusion: specifying both inferentia and trainium
   raises a clear error in step_init (mirrors @batch). Specifying
   both gpu and trainium was already enforced.
3. Wire format consistency: inferentia is popped from
   self.attributes after collapsing into trainium, so the runtime
   CLI / argo / airflow translation only ever sees the canonical key.
4. Cross-runtime: changes propagate through kubernetes_job,
   kubernetes_jobsets, argo, and airflow consistently with how
   trainium was already plumbed.
5. Validation: efa value validated as positive integer (mirrors
   trainium and tmpfs_size validation patterns in the same file).

Tests

• Unit tests added/updated
• Manual reproduction provided above
• Smoke-tested end-to-end on a real EKS cluster with Neuron and
  EFA device plugins. Pod spec contains the right resource limits;
  NCCL via aws-ofi-nccl selects EFA as the network backend.
• CI passes — TBD (CI doesn't have AWS Trainium/EFA hardware).

Non-Goals

• Not touching @batch (already has these parameters).
• Not adding a --inferentia CLI flag — inferentia is purely a
  decorator-time convenience that resolves to trainium before any
  CLI invocation, mirroring @batch's CLI which only exposes the
  canonical name (--inferentia for batch since inferentia is
  canonical there; --trainium for k8s since trainium is canonical
  here).
• Not adding NCCL/libfabric environment-variable defaults
  (FI_PROVIDER, FI_EFA_USE_DEVICE_RDMA). Users set those via
  @environment for now; auto-injection is a separate ergonomics PR.
• Not opining on which Trainium/Inferentia instance type a user
  should target — that's a cluster-side concern (instance allowlist
  • AMI selection on the EKS managed nodegroup side).

AI Tool Usage

• AI tools were used (Anthropic Claude - research on AWS DLC tag
  selection, Karpenter EFA NIC layout prior art, and drafting this
  PR description). All generated code reviewed, understood, and
  tested end-to-end on a live Outerbounds cluster.

[greptile-apps]

Greptile Summary

This PR brings @kubernetes to parity with @batch by adding trainium, inferentia, and efa decorator parameters, mapping to the aws.amazon.com/neuron and vpc.amazonaws.com/efa Kubernetes resources managed by the AWS Neuron and EFA device plugins respectively. The feature is threaded consistently through kubernetes_job, kubernetes_jobsets, kubernetes_cli, Argo, and Airflow runtimes.

• trainium=N requests N Neuron devices and automatically injects an aws.amazon.com/neuron:NoSchedule toleration; inferentia=N is a decorator-time alias that collapses to trainium in step_init before any CLI/runtime translation occurs.
• efa=N adds the vpc.amazonaws.com/efa resource limit (no toleration needed, as EFA nodes are not tainted by the device plugin).
• Both new parameters are validated as positive integers; specifying both inferentia and trainium, or both trainium and gpu, raises a clear KubernetesException.

Confidence Score: 5/5

Safe to merge — new parameters default to None, all existing paths are unaffected, and the feature has been smoke-tested end-to-end on a live EKS cluster.

The change is additive: all new attributes default to None and resource-limit emission is gated on non-None values, so no existing flow is disturbed. Validation (positive integer, mutual exclusion with gpu/trainium/inferentia) is consistent with the existing gpu pattern. The inferentia→trainium alias is safely handled before CLI serialisation — the runtime.py _options helper already filters None values, so the residual inferentia=None attribute causes no CLI errors. Neuron toleration injection follows the same structure used for gpu across all four runtimes.

No files require special attention.

Important Files Changed

Filename	Overview
metaflow/plugins/kubernetes/kubernetes_decorator.py	Adds trainium/inferentia/efa to decorator attributes, with inferentia→trainium aliasing, mutual-exclusion checks, and positive-integer validation for both trainium and efa.
metaflow/plugins/kubernetes/kubernetes_job.py	Appends aws.amazon.com/neuron and vpc.amazonaws.com/efa resource limits and the Neuron NoSchedule toleration, consistent with the existing gpu pattern.
metaflow/plugins/kubernetes/kubernetes_jobsets.py	Same trainium/efa resource limit and Neuron toleration additions as kubernetes_job.py, applied to the JobSet spec path.
metaflow/plugins/kubernetes/kubernetes_cli.py	Adds --trainium and --efa click options and threads them through the step() function to the Kubernetes execute call.
metaflow/plugins/kubernetes/kubernetes.py	Adds trainium and efa parameters to both execute method signatures and passes them through to KubernetesJob/JobSetSpec constructors.
metaflow/plugins/argo/argo_workflows.py	Passes trainium/efa through to the jobset parallel path, adds Neuron resource limits and toleration in the non-parallel container spec, consistent with existing gpu handling.
metaflow/plugins/airflow/airflow.py	Adds Neuron and EFA resource limits to the resource dict, and prepends the Neuron NoSchedule toleration to k8s_operator_args (note: Job/JobSets append instead of prepend, purely cosmetic difference).

_{Reviews (4): Last reviewed commit: "Make @kubernetes inferentia-trainium ali..." | Re-trigger Greptile}

[codecov]

Welcome to Codecov 🎉

Once you merge this PR into your default branch, you're all set! Codecov will compare coverage reports and display results in all future pull requests.

Thanks for integrating Codecov - We've got you covered ☂️

[saikonen]

No issues with the changes itself, but a question on the overall UX. Is the main goal of this feature added convenience? From what I can tell, everything seems to already be achievable with @kubernetes(tolerations=) if I'm not mistaken.

It also seems a bit of a departure for the @kubernetes decorator to implement provider-specific attributes.

[emattia]

No issues with the changes itself, but a question on the overall UX. Is the main goal of this feature added convenience? From what I can tell, everything seems to already be achievable with @kubernetes(tolerations=) if I'm not mistaken.

It also seems a bit of a departure for the @kubernetes decorator to implement provider-specific attributes.

Right, the tolerations is what needs to happen, the @kubernetes part is to make it easy for users who don't care about the backend machine. I agree it is a departure, my logic was to make the minimal change that is in parity with the existing Metaflow API on the @batch side. Ideally users don't have to set tolerations in @kubernetes imo, as many of our end users don't know what that word means. Do you have thoughts on a better approach?

[npow]

this should be greater than zero? trainium and gpu=0 should work right