Repository navigation
fix(reconciler): use the MultiKueue worker cluster for MLRuns - #1585
mattwangg-google wants to merge 3 commits into
Conversation
There was a problem hiding this comment.
Code Review
This pull request updates the JobSet reconciler to support MultiKueue execution by resolving the remote execution cluster (project, location, and cluster name) from the JobSet's Workload and MultiKueueCluster labels, ensuring that MLRun tracking records reference the correct executing cluster. Feedback on these changes highlights a compilation issue where existing tests in mlrun_sync_test.go need to be updated to match the new syncMLRun signature. Additionally, it is recommended to optimize the workload lookup in resolveExecutionCluster by registering a field indexer rather than listing and iterating over all workloads in the namespace.
|
If this is approved without comments, I'm okay to merge this directly if it can make it to this week's cut. |
|
Adding @VLukiantsev to reviewers because he has more knowledge our multikueue setup than me. |
| func (r *Reconciler) resolveExecutionCluster(ctx context.Context, jobSet *v1alpha2.JobSet) (project, location, clusterName string) { | ||
| project, location, clusterName = r.Project, r.Location, r.ClusterName | ||
|
|
||
| // Kueue labels each Workload with its owning job's UID, so this lists only the | ||
| // JobSet's Workload rather than every Workload in the namespace. | ||
| var wlList v1beta1.WorkloadList | ||
| if err := r.List(logging.Strip(ctx), &wlList, client.InNamespace(jobSet.Namespace), client.MatchingLabels{kueueJobUIDLabel: string(jobSet.UID)}); err != nil { | ||
| logging.Warn(ctx, "resolveExecutionCluster: failed to list workloads; using manager cluster", "error", err) | ||
| return project, location, clusterName | ||
| } | ||
| var wl *v1beta1.Workload | ||
| for i := range wlList.Items { | ||
| for _, or := range wlList.Items[i].OwnerReferences { | ||
| if or.Kind == "JobSet" && or.Name == jobSet.Name { | ||
| wl = &wlList.Items[i] | ||
| break | ||
| } | ||
| } | ||
| if wl != nil { | ||
| break | ||
| } | ||
| } | ||
| if wl == nil || wl.Status.ClusterName == nil || *wl.Status.ClusterName == "" { | ||
| return project, location, clusterName // not MultiKueue-dispatched | ||
| } | ||
|
|
||
| var mc v1beta1.MultiKueueCluster | ||
| if err := r.Get(logging.Strip(ctx), types.NamespacedName{Name: *wl.Status.ClusterName}, &mc); err != nil { | ||
| logging.Warn(ctx, "resolveExecutionCluster: failed to get MultiKueueCluster; using manager cluster", "cluster", *wl.Status.ClusterName, "error", err) | ||
| return project, location, clusterName | ||
| } | ||
| if mc.Labels != nil { | ||
| if v := mc.Labels["xmc.google.com/project"]; v != "" { | ||
| project = v | ||
| } | ||
| if v := mc.Labels["xmc.google.com/location"]; v != "" { | ||
| location = v | ||
| } | ||
| if v := mc.Labels["xmc.google.com/cluster-name"]; v != "" { | ||
| clusterName = v | ||
| } | ||
| } | ||
| return project, location, clusterName | ||
| } |
There was a problem hiding this comment.
[Priority: Blocker] cmd/reconciler/jobset/reconciler.go:418-461
Observation: resolveExecutionCluster silently falls back to the manager cluster on List/Get errors, before Kueue creates/admits the Workload (wl == nil or wl.Status.Admission == nil), or when MultiKueueCluster labels are partially set—which permanently caches a wrong manager-cluster MLRun URL in r.mlRunURLCache and ESS.
Risk / Rule: Cross-Component Parity (cmd/reconciler/workload/converter.go:84-108), Error Handling (golang.md#3-error-handling--formatting), and Function Length (core_engineering.md#1-srp — extract findJobSetWorkload to stay under 30 lines).
Suggestion:
| func (r *Reconciler) resolveExecutionCluster(ctx context.Context, jobSet *v1alpha2.JobSet) (project, location, clusterName string) { | |
| project, location, clusterName = r.Project, r.Location, r.ClusterName | |
| // Kueue labels each Workload with its owning job's UID, so this lists only the | |
| // JobSet's Workload rather than every Workload in the namespace. | |
| var wlList v1beta1.WorkloadList | |
| if err := r.List(logging.Strip(ctx), &wlList, client.InNamespace(jobSet.Namespace), client.MatchingLabels{kueueJobUIDLabel: string(jobSet.UID)}); err != nil { | |
| logging.Warn(ctx, "resolveExecutionCluster: failed to list workloads; using manager cluster", "error", err) | |
| return project, location, clusterName | |
| } | |
| var wl *v1beta1.Workload | |
| for i := range wlList.Items { | |
| for _, or := range wlList.Items[i].OwnerReferences { | |
| if or.Kind == "JobSet" && or.Name == jobSet.Name { | |
| wl = &wlList.Items[i] | |
| break | |
| } | |
| } | |
| if wl != nil { | |
| break | |
| } | |
| } | |
| if wl == nil || wl.Status.ClusterName == nil || *wl.Status.ClusterName == "" { | |
| return project, location, clusterName // not MultiKueue-dispatched | |
| } | |
| var mc v1beta1.MultiKueueCluster | |
| if err := r.Get(logging.Strip(ctx), types.NamespacedName{Name: *wl.Status.ClusterName}, &mc); err != nil { | |
| logging.Warn(ctx, "resolveExecutionCluster: failed to get MultiKueueCluster; using manager cluster", "cluster", *wl.Status.ClusterName, "error", err) | |
| return project, location, clusterName | |
| } | |
| if mc.Labels != nil { | |
| if v := mc.Labels["xmc.google.com/project"]; v != "" { | |
| project = v | |
| } | |
| if v := mc.Labels["xmc.google.com/location"]; v != "" { | |
| location = v | |
| } | |
| if v := mc.Labels["xmc.google.com/cluster-name"]; v != "" { | |
| clusterName = v | |
| } | |
| } | |
| return project, location, clusterName | |
| } | |
| func (r *Reconciler) resolveExecutionCluster(ctx context.Context, jobSet *v1alpha2.JobSet) (project, location, clusterName string, ok bool) { | |
| if r.Client == nil || jobSet == nil { | |
| return r.Project, r.Location, r.ClusterName, true | |
| } | |
| var wlList v1beta1.WorkloadList | |
| if err := r.List(logging.Strip(ctx), &wlList, client.InNamespace(jobSet.Namespace), client.MatchingLabels{kueueJobUIDLabel: string(jobSet.UID)}); err != nil { | |
| logging.Warn(ctx, "resolveExecutionCluster: failed to list workloads; skipping MLRun sync", "error", err) | |
| return "", "", "", false | |
| } | |
| wl := findJobSetWorkload(wlList.Items, jobSet.Name) | |
| if wl == nil || wl.Status.ClusterName == nil || *wl.Status.ClusterName == "" { | |
| if (wl == nil && jobSet.Labels["kueue.x-k8s.io/queue-name"] != "") || (wl != nil && wl.Status.Admission == nil && jobSet.Spec.Suspend != nil && *jobSet.Spec.Suspend) { | |
| return "", "", "", false | |
| } | |
| return r.Project, r.Location, r.ClusterName, true | |
| } | |
| var mc v1beta1.MultiKueueCluster | |
| if err := r.Get(logging.Strip(ctx), types.NamespacedName{Name: *wl.Status.ClusterName}, &mc); err != nil { | |
| logging.Warn(ctx, "resolveExecutionCluster: failed to get MultiKueueCluster; skipping MLRun sync", "cluster", *wl.Status.ClusterName, "error", err) | |
| return "", "", "", false | |
| } | |
| project = mc.Labels["xmc.google.com/project"] | |
| location = mc.Labels["xmc.google.com/location"] | |
| clusterName = mc.Labels["xmc.google.com/cluster-name"] | |
| if project == "" || location == "" || clusterName == "" { | |
| logging.Warn(ctx, "resolveExecutionCluster: MultiKueueCluster is missing XMC cluster labels; skipping MLRun sync", "cluster", *wl.Status.ClusterName) | |
| return "", "", "", false | |
| } | |
| return project, location, clusterName, true | |
| } | |
| func findJobSetWorkload(workloads []v1beta1.Workload, jobSetName string) *v1beta1.Workload { | |
| for i := range workloads { | |
| for _, or := range workloads[i].OwnerReferences { | |
| if or.Kind == "JobSet" && or.Name == jobSetName { | |
| return &workloads[i] | |
| } | |
| } | |
| } | |
| return nil | |
| } |
| project, location, clusterName := r.resolveExecutionCluster(ctx, jobSet) | ||
|
|
||
| mlRunURL := r.syncMLRun(ctx, jobSet, project, location, clusterName, xid, wid, eid) |
There was a problem hiding this comment.
[Priority: Actionable] cmd/reconciler/jobset/reconciler.go:399-401
Observation: r.resolveExecutionCluster(ctx, jobSet) runs WorkloadList and MultiKueueCluster lookups on every Reconcile loop of every JobSet, even though the result is only needed once when ensureMLRunCreated creates an MLRun (and never when ManagedXprof is disabled or the URL is already cached).
Risk / Rule: Reconcile Hot-Path Performance (core_engineering.md#1-srp — defer one-time resolution to ensureMLRunCreated to avoid redundant K8s lookups across high-frequency reconcile loops).
Suggestion:
| project, location, clusterName := r.resolveExecutionCluster(ctx, jobSet) | |
| mlRunURL := r.syncMLRun(ctx, jobSet, project, location, clusterName, xid, wid, eid) | |
| mlRunURL := r.syncMLRun(ctx, jobSet, xid, wid, eid) |
| // syncMLRun resolves ManagedXprof metadata for a JobSet and ensures its MLRun | ||
| // tracking record is created or reused when ManagedXprof is enabled. | ||
| func (r *Reconciler) syncMLRun(ctx context.Context, jobSet *v1alpha2.JobSet, xid, wid, eid int64) string { | ||
| func (r *Reconciler) syncMLRun(ctx context.Context, jobSet *v1alpha2.JobSet, project, location, clusterName string, xid, wid, eid int64) string { |
There was a problem hiding this comment.
[Priority: Actionable] cmd/reconciler/jobset/mlrun_sync.go:92
Observation: syncMLRun now takes 8 positional parameters and ensureMLRunCreated (line 152) takes 11 (3 strings + 3 int64s + 3 strings), and syncMLRun only passes project, location, clusterName through to ensureMLRunCreated without using them.
Risk / Rule: No Unused/Pass-Through Parameters & Function Readability (core_engineering.md#1-srp, golang.md#1-simplicity--readability).
Suggestion:
Resolve the execution cluster inside ensureMLRunCreated and pass mxEnvs (ManagedXprofEnvs) or a parameter struct instead of 11 positional arguments:
| func (r *Reconciler) syncMLRun(ctx context.Context, jobSet *v1alpha2.JobSet, project, location, clusterName string, xid, wid, eid int64) string { | |
| func (r *Reconciler) syncMLRun(ctx context.Context, jobSet *v1alpha2.JobSet, xid, wid, eid int64) string { |
| ClusterName: "cluster-a", | ||
| MLRunClient: tc.client, | ||
| } | ||
| gotURL, gotConfirmed := r.ensureMLRunCreated(context.Background(), tc.jobSet, 10, 20, 30, "display-name", "group-abc", "gs://bucket/folder") | ||
| gotURL, gotConfirmed := r.ensureMLRunCreated(context.Background(), tc.jobSet, "test-proj", "us-central1", "cluster-a", 10, 20, 30, "display-name", "group-abc", "gs://bucket/folder") |
There was a problem hiding this comment.
[Priority: Actionable] cmd/reconciler/jobset/mlrun_sync_test.go:170-173
Observation: All updated test calls (lines 170-173, 470-490, 530-543) pass the exact same "test-proj", "us-central1", "cluster-a" values as r.Project, r.Location, and r.ClusterName, so the tests would still pass even if ensureMLRunCreated ignored the execution cluster arguments.
Risk / Rule: Strict Testing Coverage (core_engineering.md#5-strict-testing-coverage--pr-granularity).
Suggestion:
| ClusterName: "cluster-a", | |
| MLRunClient: tc.client, | |
| } | |
| gotURL, gotConfirmed := r.ensureMLRunCreated(context.Background(), tc.jobSet, 10, 20, 30, "display-name", "group-abc", "gs://bucket/folder") | |
| gotURL, gotConfirmed := r.ensureMLRunCreated(context.Background(), tc.jobSet, "test-proj", "us-central1", "cluster-a", 10, 20, 30, "display-name", "group-abc", "gs://bucket/folder") | |
| ClusterName: "mgr-cluster", | |
| MLRunClient: tc.client, | |
| } | |
| gotURL, gotConfirmed := r.ensureMLRunCreated(context.Background(), tc.jobSet, "test-proj", "us-central1", "cluster-a", 10, 20, 30, "display-name", "group-abc", "gs://bucket/folder") |
| tests := []struct { | ||
| name string | ||
| objs []client.Object | ||
| want []string | ||
| }{ | ||
| {name: "no workload", want: []string{"mgr-proj", "us-central1", "mgr-gke"}}, | ||
| {name: "workload not dispatched", objs: []client.Object{workload(nil)}, want: []string{"mgr-proj", "us-central1", "mgr-gke"}}, | ||
| {name: "workload with empty cluster name", objs: []client.Object{workload(&empty)}, want: []string{"mgr-proj", "us-central1", "mgr-gke"}}, | ||
| {name: "MultiKueueCluster not found", objs: []client.Object{workload(&worker)}, want: []string{"mgr-proj", "us-central1", "mgr-gke"}}, | ||
| {name: "dispatched to worker", objs: []client.Object{workload(&worker), workerCluster}, want: []string{"worker-proj", "us-east4", "worker-gke"}}, | ||
| } |
There was a problem hiding this comment.
[Priority: Actionable] cmd/reconciler/jobset/reconciler_test.go:998-1008
Observation: TestResolveExecutionCluster is missing table cases for: (1) a Workload with non-matching OwnerReferences, (2) a MultiKueueCluster with missing/partial xmc.google.com/* labels, and (3) a WorkloadList List error.
Risk / Rule: Table-Driven Testing Coverage (core_engineering.md#5-strict-testing-coverage--pr-granularity, golang.md#5-table-driven-testing--declarative-assertions).
Suggestion:
Add table cases covering non-matching OwnerReferences, missing/partial MultiKueueCluster labels, and a List error.
|
|
||
| mlRunID := ConstructMLRunID(r.ClusterName, jobSet.Namespace, "JobSet", jobSet.Name, jobSet.CreationTimestamp.Time) | ||
| clusterPath := resolveClusterPath(r.ClusterName, r.Project, r.Location) | ||
| mlRunID := ConstructMLRunID(clusterName, jobSet.Namespace, "JobSet", jobSet.Name, jobSet.CreationTimestamp.Time) |
There was a problem hiding this comment.
[Priority: Clarity Question — potential Blocker]
Under MultiKueue, Kueue creates the worker JobSet with api.CloneObjectMetaForCreation (kueue v0.18.2 pkg/controller/jobs/jobset/jobset_multikueue_adapter.go:64), which copies only name, namespace, labels and annotations. So the worker JobSet's creationTimestamp is later than the manager's.
This ID (and workloadDetails.createTime in buildCreateMLRunOptions) now combines the worker cluster with the manager JobSet's creation time, a pair that matches neither JobSet. Managed ML Diagnostics on the worker may derive the run ID from its own JobSet (the SDK format mirrored in ConstructMLRunRawIdentifier: cluster_ns_kind_name_YYYYMMDD-HHmmss). If it does, the two IDs will differ whenever dispatch takes more than 1s, and the one-MLRun-per-JobSet conflict path won't reconcile them.
Could you confirm on a MultiKueue setup whether you end up with a single MLRun (a 409 or the "Updating existing MLRun" log) or two? If two, we'd need the worker JobSet's timestamp, or to drop the timestamp segment for MultiKueue-dispatched JobSets.
| mlRunID := ConstructMLRunID(r.ClusterName, jobSet.Namespace, "JobSet", jobSet.Name, jobSet.CreationTimestamp.Time) | ||
| clusterPath := resolveClusterPath(r.ClusterName, r.Project, r.Location) | ||
| mlRunID := ConstructMLRunID(clusterName, jobSet.Namespace, "JobSet", jobSet.Name, jobSet.CreationTimestamp.Time) | ||
| clusterPath := resolveClusterPath(clusterName, project, location) |
There was a problem hiding this comment.
[Priority: Clarity Question]
The resolved worker project/location only end up in workloadDetails.cluster. The MLRun itself is still created under the manager's MLRunProject/MLRunLocation, because the client is built once in cmd/reconciler/manager.go:219-228.
For a worker in a different region or project (the test uses us-east4 vs. us-central1), Managed ML Diagnostics on the worker will likely file its run under the worker's location. In that case the duplicate check never fires, and the XMC link opens a run with no profiles. Alternatively, the API may reject a cluster path from another location.
Was a cross-region or cross-project worker tested? If not, could you document a same-project/same-location requirement, or route the MLRun client by the resolved location?
| // manager. Resolve it from the JobSet's Workload -> MultiKueueCluster labels; | ||
| // this falls back to the manager values when the JobSet is not dispatched via | ||
| // MultiKueue. The update keeps the manager values: they populate | ||
| // metadata.managementCluster, and links use metadata.executingCluster. |
There was a problem hiding this comment.
[Priority: Actionable]
"links use metadata.executingCluster" is only partly accurate. ExecutingCluster is set only by the Workload reconciler (cmd/reconciler/workload/converter.go:51). The failure debug link built by this update (ConvertJobSetToWorkUnitExecutionUpdate → generateStructuredMessage → buildDebugLink) still uses the manager values passed a few lines below. So under MultiKueue it opens the manager's JobSet page, which has no pods.
Could you fix the comment? Passing the resolved values to the debug link would also help; a follow-up is fine for that.
| // metadata.managementCluster, and links use metadata.executingCluster. | |
| // MultiKueue. The update keeps the manager values for metadata.managementCluster; | |
| // metadata.executingCluster is set by the Workload reconciler. Note: the failure | |
| // debug link built from these manager values still targets the manager cluster. |
| return project, location, clusterName | ||
| } | ||
| if mc.Labels != nil { | ||
| if v := mc.Labels["xmc.google.com/project"]; v != "" { |
There was a problem hiding this comment.
[Priority: Nit]
The xmc.google.com/{project,location,cluster-name} MultiKueueCluster label keys are now string literals in two reconcilers (also cmd/reconciler/workload/converter.go:100-102). They're a cross-component contract, written by Helm kueue-dev/multikueue.yaml and deploy/addons/multikueue.
Could you add constants next to the existing xmc.google.com/* keys in internal/kubernetes/common/constants and share one MultiKueueCluster → ClusterInfo helper? That way the MLRun cluster and executing_cluster can't silently drift apart.
This PR handles the configuration of the MLRun to target the worker cluster if MultiKueue is used.
This is achieved by adding the project, location, and cluster name as configurable parameters. The default is to use the existing ones known by the reconciler, but if MultiKueue is used, then the worker cluster is retrieved to provide the right data.
This helps in situations where MultiKueue is used.
It shouldn't affect current default behavior, as if MultiKueue isn't used, then it should fall back to the defaults.
Tested with unit tests and MTC team using XMC.