[2.25] Backport of the kueue integration and upgrade tests for the workbenches component

[jstourac]

Pull Request

Summary

Backport into the 2.25 branch of the:

• test(workbenches): add Kueue upgrade tests for notebook scheduling #2063
• test(workbenches): add Kueue integration tests for notebook scheduling #2062
  with some necessary changes to adapt to the 2.25 product reality and some extra specific modifications. Two of the commits from here may be actually migrated to a newer branches too eventually.

Basically this is backport of the #2108 with those extra adaptations for 2.25 world.

Related Issues

https://redhat.atlassian.net/browse/RHOAIENG-73662

Please review and indicate how it has been tested

• Locally
• Jenkins

Additional Requirements

• If this PR introduces a new test image, did you create a PR to mirror it in disconnected environment?
• If this PR introduces new marker(s)/adds a new component, was relevant ticket created to update relevant Jenkins job?

Summary by CodeRabbit

• Tests
  • Added integration coverage for queue-based notebook scheduling, including admission control, resource limits, and stop/start behavior.
  • Added upgrade coverage verifying that running and stopped notebooks retain their state, queue assignments, labels, and workloads across upgrades.
  • Added validation for creating and admitting new queued notebooks after an upgrade.
  • Improved notebook readiness checks with clearer timeout diagnostics and consistent handling across standard and upgrade scenarios.

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: 62e689cf-a120-4d5b-b715-8e939688213a

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

📝 Walkthrough

Walkthrough

Changes

Kueue notebook integration and upgrade coverage

Layer / File(s)	Summary
Shared Kueue and notebook test support utilities/kueue_utils.py, tests/workbenches/notebooks_server/controller/utils.py, tests/workbenches/notebooks_server/controller/conftest.py	Adds Kueue resource models, HardwareProfile support, configurable notebook labels and resources, pod readiness handling, and Kueue resource fixtures.
Kueue notebook scheduling scenarios tests/workbenches/notebooks_server/controller/test_kueue_integration.py	Tests admission, quota-based gating, queue labels, resource injection, and stop/start re-admission.
Kueue upgrade state and fixtures tests/workbenches/notebooks_server/controller/upgrade/kueue_constants.py, tests/workbenches/notebooks_server/controller/upgrade/conftest.py	Adds upgrade resource fixtures, stopped-notebook handling, baseline capture, and shared readiness-helper usage.
Kueue upgrade validation tests/workbenches/notebooks_server/controller/upgrade/test_upgrade_kueue.py	Validates resource preservation, stopped-notebook state, workload admission, and creation of new Kueue-managed notebooks after upgrade.

Estimated code review effort: 4 (Complex) | ~60 minutes

Possibly related PRs

• opendatahub-io/opendatahub-tests#2062: Extends overlapping Kueue integration fixtures, utilities, and upgrade coverage.
• opendatahub-io/opendatahub-tests#1962: Overlaps the notebook pod readiness timeout handling.
• opendatahub-io/opendatahub-tests#2063: Covers overlapping Kueue integration and upgrade test files.

Suggested reviewers: andyatmiami, harshad16, jiridanek

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies the backport of Kueue integration and upgrade tests for the Workbenches component.
Description check	✅ Passed	The description includes the summary, related issue, testing status, and additional requirement checklist required by the template.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[jstourac]

@coderabbitai full review

[coderabbitai]

✅ Action performed

Full review finished.

[github-actions]

The following are automatically added/executed:

• PR size label.
• Run pre-commit
• Run tox
• Add PR author as the PR assignee
• Build image based on the PR

Available user actions:

• To mark a PR as WIP, add /wip in a comment. To remove it from the PR comment /wip cancel to the PR.
• To block merging of a PR, add /hold in a comment. To un-block merging of PR comment /hold cancel.
• To mark a PR as approved, add /lgtm in a comment. To remove, add /lgtm cancel.
  lgtm label removed on each new commit push.
• To mark PR as verified comment /verified to the PR, to un-verify comment /verified cancel to the PR.
  verified label removed on each new commit push.
• To Cherry-pick a merged PR /cherry-pick <target_branch_name> to the PR. If <target_branch_name> is valid,
  and the current PR is merged, a cherry-picked PR would be created and linked to the current PR.
• To build and push image to quay, add /build-push-pr-image in a comment. This would create an image with tag
  pr-<pr_number> to quay repository. This image tag, however would be deleted on PR merge or close action.

Supported labels

{'/lgtm', '/build-push-pr-image', '/cherry-pick', '/wip', '/verified', '/hold'}

[coderabbitai]

Actionable comments posted: 2

🧹 Nitpick comments (7)

utilities/kueue_utils.py (1)

9-10: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Kueue label keys are defined in three places. utilities/constants.py already holds Labels.Kueue.QUEUE_NAME and Labels.Kueue.MANAGED. This PR adds a second copy in utilities/kueue_utils.py and a third private copy in the integration test. The shared root cause is the absence of one source of truth, so a future key change can leave the copies out of sync and turn label assertions into no-ops.

• utilities/kueue_utils.py#L9-L10: reference or re-export Labels.Kueue.QUEUE_NAME instead of redefining KUEUE_QUEUE_NAME_LABEL, and confirm whether KUEUE_MANAGED_LABEL should equal Labels.Kueue.MANAGED or is a distinct pod-level key.
• tests/workbenches/notebooks_server/controller/test_kueue_integration.py#L43-L44: delete _KUEUE_QUEUE_NAME_LABEL and _KUBEFLOW_STOPPED_ANNOTATION, and import KUEUE_QUEUE_NAME_LABEL from utilities.kueue_utils and KUBEFLOW_STOPPED_ANNOTATION from tests/workbenches/notebooks_server/controller/utils.py.

As per coding guidelines: "Avoid code duplication by creating meaningful utilities".

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@utilities/kueue_utils.py` around lines 9 - 10, The Kueue label and annotation
keys are duplicated across utilities and integration tests. In
utilities/kueue_utils.py lines 9-10, reference or re-export
Labels.Kueue.QUEUE_NAME for KUEUE_QUEUE_NAME_LABEL and verify whether
KUEUE_MANAGED_LABEL should use Labels.Kueue.MANAGED or remain a distinct
pod-level key; in
tests/workbenches/notebooks_server/controller/test_kueue_integration.py lines
43-44, remove the private _KUEUE_QUEUE_NAME_LABEL and
_KUBEFLOW_STOPPED_ANNOTATION definitions and import the shared
KUEUE_QUEUE_NAME_LABEL and KUBEFLOW_STOPPED_ANNOTATION symbols.

Source: Coding guidelines

tests/workbenches/notebooks_server/controller/conftest.py (2)

209-210: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Rename this fixture to a noun.

kueue_statefulset_framework_check reads as an action. The repository requires noun fixture names. Consider kueue_statefulset_framework or kueue_statefulset_integration.

As per coding guidelines: "Fixture names MUST be nouns (e.g., storage_secret, not create_secret)".

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/workbenches/notebooks_server/controller/conftest.py` around lines 209 -
210, Rename the kueue_statefulset_framework_check fixture to a noun-based name,
preferably kueue_statefulset_framework, and update every reference to the
fixture consistently without changing its behavior.

Source: Coding guidelines

---

98-163: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Reuse wait_for_notebook_pod_ready in this fixture.

This file already imports helpers from tests/workbenches/notebooks_server/controller/utils.py. That module now provides wait_for_notebook_pod_ready (lines 281-310), which performs the same wait, diagnostic collection, and assertion split. Keeping both copies means future fixes must land twice.

Import the helper and pass the desired timeout.

♻️ Proposed refactor

-    _ERR_POD_NOT_READY = (
-        "Pod '{pod_name}-0' failed to reach Ready state within 10 minutes.\n"
-        "Pod Phase: {pod_phase}\n"
-        "Original Error: {original_error}\n"
-        "Pod information collected to must-gather directory for debugging."
-    )
-    _ERR_POD_NOT_CREATED = "Pod '{pod_name}-0' was not created. Check notebook controller logs."
-
     notebook_pod = Pod(
         client=unprivileged_client,
         namespace=default_notebook.namespace,
         name=f"{default_notebook.name}-0",
     )
 
-    try:
-        notebook_pod.wait()
-        ...
-            raise AssertionError(_ERR_POD_NOT_CREATED.format(pod_name=default_notebook.name)) from e
+    wait_for_notebook_pod_ready(
+        notebook_pod=notebook_pod,
+        context="Default notebook",
+        timeout=Timeout.TIMEOUT_10MIN,
+    )
 
     return notebook_pod

As per coding guidelines: "Avoid code duplication by creating meaningful utilities".

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/workbenches/notebooks_server/controller/conftest.py` around lines 98 -
163, Replace the duplicated wait, diagnostics, and assertion logic in the
notebook_pod fixture with the existing wait_for_notebook_pod_ready helper from
the controller utilities module. Import that helper and invoke it with
notebook_pod and the desired 10-minute timeout, preserving the fixture’s
returned Ready pod behavior.

Source: Coding guidelines

tests/workbenches/notebooks_server/controller/upgrade/test_upgrade_kueue.py (2)

254-293: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Rename the class or move the pre-upgrade tests.

TestPostUpgradeKueueStopped contains test_kueue_stopped_notebook_has_zero_replicas and test_kueue_stopped_notebook_pod_absent, and both carry @pytest.mark.pre_upgrade. The class name states post-upgrade only. The neighbouring classes separate the two phases (TestPreUpgradeKueueNotebook, TestPostUpgradeKueueNotebook).

Split the pre-upgrade tests into a TestPreUpgradeKueueStopped class, or rename this class to cover both phases.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/workbenches/notebooks_server/controller/upgrade/test_upgrade_kueue.py`
around lines 254 - 293, Move the pre-upgrade tests
test_kueue_stopped_notebook_has_zero_replicas and
test_kueue_stopped_notebook_pod_absent from TestPostUpgradeKueueStopped into a
TestPreUpgradeKueueStopped class, preserving their `@pytest.mark.pre_upgrade`
markers and fixture usage; keep TestPostUpgradeKueueStopped reserved for
post-upgrade tests.

---

233-251: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Extract the Workload admission assertion.

Both blocks perform the same four steps: list Workload objects in the namespace, filter by notebook name, check the Admitted condition, and compare status.admission.clusterQueue. tests/workbenches/notebooks_server/controller/test_kueue_integration.py already defines _get_notebook_workload and _workload_is_admitted for this purpose.

Move those two helpers to a shared location and call them from both classes.

Also applies to: 429-445

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/workbenches/notebooks_server/controller/upgrade/test_upgrade_kueue.py`
around lines 233 - 251, Extract the duplicated Workload lookup and admission
checks from the upgrade test into a shared helper location, reusing the existing
`_get_notebook_workload` and `_workload_is_admitted` behavior from
`test_kueue_integration.py`. Update both test classes, including the
corresponding block at the later referenced section, to call the shared helpers
instead of listing, filtering, and asserting admission inline; preserve the
existing ClusterQueue validation.

tests/workbenches/notebooks_server/controller/upgrade/conftest.py (2)

972-989: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Move the Workload garbage-collection assertion out of teardown.

pytest.fail inside a fixture teardown reports an error against the fixture finalizer, not against a named test. The failure is hard to attribute, and it runs only when teardown_resources is true, so the check disappears in retain-resources runs. The behaviour under test here is a product guarantee: Kueue must garbage-collect Workload objects after the notebook StatefulSet is deleted.

Add a dedicated post-upgrade test that deletes the notebook and asserts the Workload cleanup. Keep the teardown limited to deletion, and log a warning if Workloads remain.

♻️ Proposed teardown simplification

         if teardown_resources:
             nb.client = admin_client
             nb.clean_up()
             try:
                 for sample in TimeoutSampler(
                     wait_timeout=60,
                     sleep=5,
                     func=lambda: list(Workload.get(client=admin_client, namespace=upgrade_kueue_namespace.name)),
                 ):
                     if not sample:
                         break
             except TimeoutExpiredError:
                 remaining = list(Workload.get(client=admin_client, namespace=upgrade_kueue_namespace.name))
-                pytest.fail(
-                    f"Kueue did not clean up {len(remaining)} Workload(s) within 60s after notebook deletion: "
-                    f"{[wl.name for wl in remaining]}. "
-                    f"This indicates Kueue is not garbage-collecting Workloads for deleted StatefulSets."
-                )
+                LOGGER.warning(
+                    f"Kueue did not clean up {len(remaining)} Workload(s) within 60s after notebook deletion: "
+                    f"{[wl.name for wl in remaining]}"
+                )

I can draft the dedicated test if you want.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/workbenches/notebooks_server/controller/upgrade/conftest.py` around
lines 972 - 989, Move the Workload garbage-collection assertion from the
teardown block in the upgrade fixture into a dedicated post-upgrade test with a
named test function. Have that test delete the notebook, poll until Workloads in
upgrade_kueue_namespace are gone, and fail the test with the existing diagnostic
when cleanup times out. Keep fixture teardown limited to deleting resources, and
replace teardown-time failure with a warning when Workloads remain.

---

1166-1166: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Use the shared timeout constant.

The sibling fixture at line 425 uses Timeout.TIMEOUT_2MIN for the same wait. Timeout is already imported.

♻️ Proposed change

-    notebook_pod.wait_deleted(timeout=120)
+    notebook_pod.wait_deleted(timeout=Timeout.TIMEOUT_2MIN)

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/workbenches/notebooks_server/controller/upgrade/conftest.py` at line
1166, Update the cleanup wait in the affected fixture to pass the shared
Timeout.TIMEOUT_2MIN constant to notebook_pod.wait_deleted instead of the
hard-coded 120 value, matching the sibling fixture’s existing timeout usage.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@tests/workbenches/notebooks_server/controller/conftest.py`:
- Around line 251-283: Initialize the Kueue configuration’s integrations section
once after loading config_data, using an empty mapping when it is absent, and
reuse that reference for reading frameworks and writing frameworks/podOptions.
Update the relevant logic around config_data and frameworks_list so missing
integrations sections are created safely; add the required Any import if the
chosen typing annotation needs it.

In `@tests/workbenches/notebooks_server/controller/utils.py`:
- Around line 40-43: Update the HardwareProfile class docstring to reference
infrastructure.opendatahub.io/v1alpha1 instead of v1, matching the apiVersion
used by its consumers.

---

Nitpick comments:
In `@tests/workbenches/notebooks_server/controller/conftest.py`:
- Around line 209-210: Rename the kueue_statefulset_framework_check fixture to a
noun-based name, preferably kueue_statefulset_framework, and update every
reference to the fixture consistently without changing its behavior.
- Around line 98-163: Replace the duplicated wait, diagnostics, and assertion
logic in the notebook_pod fixture with the existing wait_for_notebook_pod_ready
helper from the controller utilities module. Import that helper and invoke it
with notebook_pod and the desired 10-minute timeout, preserving the fixture’s
returned Ready pod behavior.

In `@tests/workbenches/notebooks_server/controller/upgrade/conftest.py`:
- Around line 972-989: Move the Workload garbage-collection assertion from the
teardown block in the upgrade fixture into a dedicated post-upgrade test with a
named test function. Have that test delete the notebook, poll until Workloads in
upgrade_kueue_namespace are gone, and fail the test with the existing diagnostic
when cleanup times out. Keep fixture teardown limited to deleting resources, and
replace teardown-time failure with a warning when Workloads remain.
- Line 1166: Update the cleanup wait in the affected fixture to pass the shared
Timeout.TIMEOUT_2MIN constant to notebook_pod.wait_deleted instead of the
hard-coded 120 value, matching the sibling fixture’s existing timeout usage.

In `@tests/workbenches/notebooks_server/controller/upgrade/test_upgrade_kueue.py`:
- Around line 254-293: Move the pre-upgrade tests
test_kueue_stopped_notebook_has_zero_replicas and
test_kueue_stopped_notebook_pod_absent from TestPostUpgradeKueueStopped into a
TestPreUpgradeKueueStopped class, preserving their `@pytest.mark.pre_upgrade`
markers and fixture usage; keep TestPostUpgradeKueueStopped reserved for
post-upgrade tests.
- Around line 233-251: Extract the duplicated Workload lookup and admission
checks from the upgrade test into a shared helper location, reusing the existing
`_get_notebook_workload` and `_workload_is_admitted` behavior from
`test_kueue_integration.py`. Update both test classes, including the
corresponding block at the later referenced section, to call the shared helpers
instead of listing, filtering, and asserting admission inline; preserve the
existing ClusterQueue validation.

In `@utilities/kueue_utils.py`:
- Around line 9-10: The Kueue label and annotation keys are duplicated across
utilities and integration tests. In utilities/kueue_utils.py lines 9-10,
reference or re-export Labels.Kueue.QUEUE_NAME for KUEUE_QUEUE_NAME_LABEL and
verify whether KUEUE_MANAGED_LABEL should use Labels.Kueue.MANAGED or remain a
distinct pod-level key; in
tests/workbenches/notebooks_server/controller/test_kueue_integration.py lines
43-44, remove the private _KUEUE_QUEUE_NAME_LABEL and
_KUBEFLOW_STOPPED_ANNOTATION definitions and import the shared
KUEUE_QUEUE_NAME_LABEL and KUBEFLOW_STOPPED_ANNOTATION symbols.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: ea32c4a3-1b25-485f-a109-9e1a3c5012d9

📥 Commits

Reviewing files that changed from the base of the PR and between 5722f22 and 0dee92e.

📒 Files selected for processing (7)

• tests/workbenches/notebooks_server/controller/conftest.py
• tests/workbenches/notebooks_server/controller/test_kueue_integration.py
• tests/workbenches/notebooks_server/controller/upgrade/conftest.py
• tests/workbenches/notebooks_server/controller/upgrade/kueue_constants.py
• tests/workbenches/notebooks_server/controller/upgrade/test_upgrade_kueue.py
• tests/workbenches/notebooks_server/controller/utils.py
• utilities/kueue_utils.py

[jstourac]

@coderabbitai full review

[coderabbitai]

✅ Action performed

Full review finished.

---

Your included review limit is currently reached under our Fair Usage Limits Policy. This review may still proceed through usage-based billing if eligible. Your next included review will be available in 33 minutes.

[github-actions]

Status of building tag 2.25: success.
Status of pushing tag 2.25 to image registry: success.
Image manifest label: success.