feat(trainer): Support NPU labels in TrainJob device by sujalshah-bit · Pull Request #336 · kubeflow/sdk

sujalshah-bit · 2026-02-27T19:45:06Z

What this PR does / why we need it:
Add support for NPU resource labels in container resource limits,
resolving the existing TODO to support additional accelerator types
(e.g., NPUs).

Detect NPU resource keys using strict /npu suffix matching
Extract and return the corresponding resource value
Fails fast when multiple NPU resource types are present
Maintains consistency with existing GPU/MIG validation behavior

Fixes NONE

Checklist:

Docs included if any changes are user facing

google-oss-prow · 2026-02-27T19:45:13Z

[APPROVALNOTIFIER] This PR is NOT APPROVED

This pull-request has been approved by:
Once this PR has been reviewed and has the lgtm label, please assign andreyvelich for approval. For more information see the Kubernetes Code Review Process.

The full list of commands accepted by this bot can be found here.

Details

Needs approval from an approver in each of these files:

OWNERS

Approvers can indicate their approval by writing /approve in a comment
Approvers can cancel approval by writing /approve cancel in a comment

github-actions · 2026-02-27T19:45:17Z

🎉 Welcome to the Kubeflow SDK! 🎉

Thanks for opening your first PR! We're happy to have you as part of our community 🚀

Here's what happens next:

If you haven't already, please check out our Contributing Guide for repo-specific guidelines and the Kubeflow Contributor Guide for general community standards
Our team will review your PR soon! cc @kubeflow/kubeflow-sdk-team

Join the community:

Slack: Join our #kubeflow-ml-experience and #kubeflow-trainer Slack channels
Meetings: Attend the Kubeflow SDK and ML Experience bi-weekly meetings

Feel free to ask questions in the comments if you need any help or clarification!
Thanks again for contributing to Kubeflow! 🙏

Copilot

Pull request overview

Adds NPU extended-resource detection to the Kubernetes backend so TrainJob device metadata can be derived from container resource limits when an NPU resource key is present.

Changes:

Introduces an NPU-related resource constant in constants.py.
Extends get_container_devices() to detect */npu resource keys and fail fast on multiple NPU types.
Adds unit tests covering single-NPU and multiple-NPU-key scenarios.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.

File	Description
kubeflow/trainer/constants/constants.py	Adds an NPU-related constant for resource identification.
kubeflow/trainer/backends/kubernetes/utils.py	Adds NPU resource key detection and validation in device extraction.
kubeflow/trainer/backends/kubernetes/utils_test.py	Adds test cases validating NPU detection and multi-key failure behavior.

Copilot · 2026-02-27T19:47:10Z

kubeflow/trainer/constants/constants.py

+# The label for NPU in the container resources.
+NPU_LABEL = "/npu"
+


NPU_LABEL is defined as the suffix string "/npu" (not a full resource label like the other *_LABEL constants) and is currently unused, which is likely to confuse future callers; rename it to something like NPU_RESOURCE_SUFFIX/NPU_SUFFIX (or document it explicitly as a suffix) and update usages accordingly.

Copilot · 2026-02-27T19:47:10Z

kubeflow/trainer/backends/kubernetes/utils.py

        mig_key = mig_keys[0]
        device = mig_key.split("/")[1]
        device_count = resources.limits[mig_key].actual_instance
+    elif npu_keys := [k for k in resources.limits if k.endswith("/npu")]:


The NPU detection uses a hard-coded string literal "/npu" even though an NPU constant was added in constants; use the constant here (and/or rename it to a suffix constant) so the matching rule is defined in one place and can’t drift.

Suggested change

elif npu_keys := [k for k in resources.limits if k.endswith("/npu")]:

elif npu_keys := [k for k in resources.limits if k.endswith(constants.NPU_SUFFIX)]:

Implement support for NPU resource labels in resource limit validation, resolving the existing TODO to support additional accelerator types. Signed-off-by: Sujal Shah <sujalshah28092004@gmail.com>

sujalshah-bit · 2026-03-12T04:09:45Z

@andreyvelich and @kramaranya, gentle reminder for the review :)

andreyvelich · 2026-02-28T01:02:08Z

kubeflow/trainer/backends/kubernetes/utils.py

    # Potentially, we can use the trainer.kubeflow.org/device label from the runtime or
    # node types.
-    # TODO (andreyvelich): Support other resource labels (e.g. NPUs).
+    # TODO (andreyvelich): Support other resource labels.


You can remove this TODO statement.

Suggested change

# TODO (andreyvelich): Support other resource labels.

andreyvelich · 2026-03-12T12:19:28Z

kubeflow/trainer/backends/kubernetes/utils.py

+        if len(npu_keys) > 1:
+            raise ValueError(f"Multiple NPU resource types are not supported yet: {npu_keys}")


Do we need this?

andreyvelich · 2026-03-12T12:20:28Z

kubeflow/trainer/backends/kubernetes/utils.py

+        if len(npu_keys) > 1:
+            raise ValueError(f"Multiple NPU resource types are not supported yet: {npu_keys}")
+
+        npu_key = npu_keys[0]


Why are you taking the first item? I thought NPU resources are set similarly to TPU and GPU

andreyvelich · 2026-03-12T12:20:40Z

kubeflow/trainer/backends/kubernetes/utils_test.py

+            config={
+                "resources": models.IoK8sApiCoreV1ResourceRequirements(
+                    limits={
+                        "huawei.com/npu": models.IoK8sApimachineryPkgApiResourceQuantity(2),


Remove vendors from unit tests

Suggested change

"huawei.com/npu": models.IoK8sApimachineryPkgApiResourceQuantity(2),

"example.com/npu": models.IoK8sApimachineryPkgApiResourceQuantity(2),

andreyvelich · 2026-03-12T12:20:50Z

kubeflow/trainer/backends/kubernetes/utils_test.py

+            expected_output=("npu", "2.0"),
+        ),
+        TestCase(
+            name="multiple NPU resource types are not supported",


This should not be the case.

andreyvelich · 2026-03-12T12:24:38Z

kubeflow/trainer/constants/constants.py

 TPU_LABEL = "google.com/tpu"

+# The Suffix label for NPU in the container resources.
+NPU_LABEL_SUFFIX = "/npu"


I thought the resource name is: huawei.com/Ascend910.
@danish9039 Do you know details about it?
Ref: #264

andreyvelich · 2026-03-12T12:24:59Z

/ok-to-test

Copilot AI review requested due to automatic review settings February 27, 2026 19:45

google-oss-prow bot requested review from andreyvelich, kramaranya and szaher February 27, 2026 19:45

google-oss-prow bot added the size/M label Feb 27, 2026

Copilot started reviewing on behalf of sujalshah-bit February 27, 2026 19:45 View session

Copilot AI reviewed Feb 27, 2026

View reviewed changes

feat(trainer): Support NPU labels in TrainJob device

6975af2

Implement support for NPU resource labels in resource limit validation, resolving the existing TODO to support additional accelerator types. Signed-off-by: Sujal Shah <sujalshah28092004@gmail.com>

sujalshah-bit force-pushed the support_npu_label branch from f4fa8fe to 6975af2 Compare February 27, 2026 20:07

andreyvelich reviewed Mar 12, 2026

View reviewed changes

google-oss-prow bot added the ok-to-test label Mar 12, 2026

		# The label for NPU in the container resources.
		NPU_LABEL = "/npu"

	elif npu_keys := [k for k in resources.limits if k.endswith("/npu")]:
	elif npu_keys := [k for k in resources.limits if k.endswith(constants.NPU_SUFFIX)]:

		if len(npu_keys) > 1:
		raise ValueError(f"Multiple NPU resource types are not supported yet: {npu_keys}")

	"huawei.com/npu": models.IoK8sApimachineryPkgApiResourceQuantity(2),
	"example.com/npu": models.IoK8sApimachineryPkgApiResourceQuantity(2),

Conversation

sujalshah-bit commented Feb 27, 2026

Uh oh!

google-oss-prow bot commented Feb 27, 2026

Uh oh!

github-actions bot commented Feb 27, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Copilot AI Feb 27, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Feb 27, 2026

Choose a reason for hiding this comment

Uh oh!

sujalshah-bit commented Mar 12, 2026

Uh oh!

andreyvelich Feb 28, 2026

Choose a reason for hiding this comment

Uh oh!

andreyvelich Mar 12, 2026

Choose a reason for hiding this comment

Uh oh!

andreyvelich Mar 12, 2026

Choose a reason for hiding this comment

Uh oh!

andreyvelich Mar 12, 2026

Choose a reason for hiding this comment

Uh oh!

andreyvelich Mar 12, 2026

Choose a reason for hiding this comment

Uh oh!

andreyvelich Mar 12, 2026

Choose a reason for hiding this comment

Uh oh!

andreyvelich commented Mar 12, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants