From 488a17cabbeb27f89cb3b9a281dd8bc91bca0c2b Mon Sep 17 00:00:00 2001 From: Aravind Warrier Date: Tue, 28 Jul 2026 16:23:12 -0400 Subject: [PATCH 1/4] feat(controller): Reconcile TriageRun jobs --- api/v2/application_types.go | 56 ++ api/v2/zz_generated.deepcopy.go | 64 ++ cmd/manager/main.go | 17 + .../bases/apps.wandb.com_applications.yaml | 148 ++++ config/rbac/role.yaml | 9 + config/samples/apps_v2_application.yaml | 15 +- .../templates/wandb-operator-wandb-role.yaml | 11 +- .../controller/reconciler/reconcile_v2.go | 1 + internal/controller/reconciler/triage.go | 34 + internal/controller/reconciler/triage_test.go | 53 ++ internal/controller/triagerun_controller.go | 724 ++++++++++++++++++ .../triagerun_controller_unit_test.go | 340 ++++++++ .../operator/apps.wandb.com_applications.yaml | 148 ++++ pkg/wandb/manifest/manifest.go | 17 + 14 files changed, 1635 insertions(+), 2 deletions(-) create mode 100644 internal/controller/reconciler/triage.go create mode 100644 internal/controller/reconciler/triage_test.go create mode 100644 internal/controller/triagerun_controller.go create mode 100644 internal/controller/triagerun_controller_unit_test.go diff --git a/api/v2/application_types.go b/api/v2/application_types.go index c5cf1d9c..e8a28352 100644 --- a/api/v2/application_types.go +++ b/api/v2/application_types.go @@ -67,11 +67,67 @@ type ApplicationSpec struct { Jobs []batchv1.Job `json:"jobs,omitempty"` CronJobs []batchv1.CronJob `json:"cronJobs,omitempty"` + // Triage declares the bounded diagnostic actions that may be requested for + // this application through TriageRun resources. + // +optional + Triage *ApplicationTriageSpec `json:"triage,omitempty"` + // HTTPRouteTemplate is the desired HTTPRoute spec. Nil means no HTTPRoute. // +optional HTTPRouteTemplate *HTTPRouteTemplateSpec `json:"httpRouteTemplate,omitempty"` } +// ApplicationTriageSpec contains the diagnostic actions exposed by an +// Application. The Application remains the source of runtime configuration; +// TriageRun only selects one of these actions. +type ApplicationTriageSpec struct { + // Actions maps stable action names to their execution overrides. + // +kubebuilder:validation:MinProperties=1 + // +kubebuilder:validation:MaxProperties=16 + Actions map[string]TriageActionSpec `json:"actions"` +} + +// TriageActionSpec is a compact override applied to one container from the +// Application pod template. Image, identity, environment, volumes, and +// scheduling settings are inherited from the Application. +type TriageActionSpec struct { + // ContainerName selects a container from the Application pod template. It + // may be omitted when the Application has exactly one container. + // +optional + ContainerName string `json:"containerName,omitempty"` + + // Command replaces the selected container's entrypoint when non-empty. + // +kubebuilder:validation:MinItems=1 + // +kubebuilder:validation:MaxItems=64 + // +optional + Command []string `json:"command,omitempty"` + + // Args replaces the selected container's arguments when non-empty. + // +kubebuilder:validation:MinItems=1 + // +kubebuilder:validation:MaxItems=64 + // +optional + Args []string `json:"args,omitempty"` + + // Env adds or overrides environment variables inherited from the selected + // application container. + // +kubebuilder:validation:MaxItems=128 + // +optional + Env []corev1.EnvVar `json:"env,omitempty"` + + // Resources deliberately does not inherit the parent container's resource + // requirements. When omitted, the controller applies small bounded + // defaults suitable for diagnostics. + // +optional + Resources *corev1.ResourceRequirements `json:"resources,omitempty"` + + // TimeoutSeconds is the Job execution deadline. Zero selects the controller + // default. + // +kubebuilder:validation:Minimum=1 + // +kubebuilder:validation:Maximum=3600 + // +optional + TimeoutSeconds int64 `json:"timeoutSeconds,omitempty"` +} + // HTTPRouteTemplateSpec contains the fields needed to build a Gateway API HTTPRoute. type HTTPRouteTemplateSpec struct { ParentRefs []gatewayv1.ParentReference `json:"parentRefs"` diff --git a/api/v2/zz_generated.deepcopy.go b/api/v2/zz_generated.deepcopy.go index 93c5bd13..514b095a 100644 --- a/api/v2/zz_generated.deepcopy.go +++ b/api/v2/zz_generated.deepcopy.go @@ -150,6 +150,11 @@ func (in *ApplicationSpec) DeepCopyInto(out *ApplicationSpec) { (*in)[i].DeepCopyInto(&(*out)[i]) } } + if in.Triage != nil { + in, out := &in.Triage, &out.Triage + *out = new(ApplicationTriageSpec) + (*in).DeepCopyInto(*out) + } if in.HTTPRouteTemplate != nil { in, out := &in.HTTPRouteTemplate, &out.HTTPRouteTemplate *out = new(HTTPRouteTemplateSpec) @@ -231,6 +236,28 @@ func (in *ApplicationStatus) DeepCopy() *ApplicationStatus { return out } +// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. +func (in *ApplicationTriageSpec) DeepCopyInto(out *ApplicationTriageSpec) { + *out = *in + if in.Actions != nil { + in, out := &in.Actions, &out.Actions + *out = make(map[string]TriageActionSpec, len(*in)) + for key, val := range *in { + (*out)[key] = *val.DeepCopy() + } + } +} + +// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new ApplicationTriageSpec. +func (in *ApplicationTriageSpec) DeepCopy() *ApplicationTriageSpec { + if in == nil { + return nil + } + out := new(ApplicationTriageSpec) + in.DeepCopyInto(out) + return out +} + // DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. func (in *CertManagerConfig) DeepCopyInto(out *CertManagerConfig) { *out = *in @@ -1494,6 +1521,43 @@ func (in *TriageActionReference) DeepCopy() *TriageActionReference { return out } +// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. +func (in *TriageActionSpec) DeepCopyInto(out *TriageActionSpec) { + *out = *in + if in.Command != nil { + in, out := &in.Command, &out.Command + *out = make([]string, len(*in)) + copy(*out, *in) + } + if in.Args != nil { + in, out := &in.Args, &out.Args + *out = make([]string, len(*in)) + copy(*out, *in) + } + if in.Env != nil { + in, out := &in.Env, &out.Env + *out = make([]v1.EnvVar, len(*in)) + for i := range *in { + (*in)[i].DeepCopyInto(&(*out)[i]) + } + } + if in.Resources != nil { + in, out := &in.Resources, &out.Resources + *out = new(v1.ResourceRequirements) + (*in).DeepCopyInto(*out) + } +} + +// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new TriageActionSpec. +func (in *TriageActionSpec) DeepCopy() *TriageActionSpec { + if in == nil { + return nil + } + out := new(TriageActionSpec) + in.DeepCopyInto(out) + return out +} + // DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. func (in *TriageActionStatus) DeepCopyInto(out *TriageActionStatus) { *out = *in diff --git a/cmd/manager/main.go b/cmd/manager/main.go index 7edc029a..eb91036d 100644 --- a/cmd/manager/main.go +++ b/cmd/manager/main.go @@ -40,6 +40,7 @@ import ( "github.com/wandb/operator/pkg/wandb/spec/channel/deployer" corev1 "k8s.io/api/core/v1" "k8s.io/client-go/discovery" + "k8s.io/client-go/kubernetes" "sigs.k8s.io/controller-runtime/pkg/client/config" gatewayv1 "sigs.k8s.io/gateway-api/apis/v1" @@ -339,6 +340,22 @@ func main() { os.Exit(1) } + kubernetesClient, err := kubernetes.NewForConfig(mgr.GetConfig()) + if err != nil { + setupLog.Error(err, "unable to create Kubernetes client", "controller", "TriageRun") + os.Exit(1) + } + if err = (&controller.TriageRunReconciler{ + Client: mgr.GetClient(), + Scheme: mgr.GetScheme(), + PodLogs: &controller.KubernetesTriagePodLogReader{ + CoreV1: kubernetesClient.CoreV1(), + }, + }).SetupWithManager(mgr); err != nil { + setupLog.Error(err, "unable to create controller", "controller", "TriageRun") + os.Exit(1) + } + if enableWebhooks && enableV2 { if err := webhookv2.SetupApplicationWebhookWithManager(mgr); err != nil { setupLog.Error(err, "unable to create webhook", "webhook", "Application") diff --git a/config/crd/bases/apps.wandb.com_applications.yaml b/config/crd/bases/apps.wandb.com_applications.yaml index 67be4bc9..5cd5c9dd 100644 --- a/config/crd/bases/apps.wandb.com_applications.yaml +++ b/config/crd/bases/apps.wandb.com_applications.yaml @@ -12801,6 +12801,154 @@ spec: type: type: string type: object + triage: + properties: + actions: + additionalProperties: + properties: + args: + items: + type: string + maxItems: 64 + minItems: 1 + type: array + command: + items: + type: string + maxItems: 64 + minItems: 1 + type: array + containerName: + type: string + env: + items: + properties: + name: + type: string + value: + type: string + valueFrom: + properties: + configMapKeyRef: + properties: + key: + type: string + name: + default: "" + type: string + optional: + type: boolean + required: + - key + type: object + x-kubernetes-map-type: atomic + fieldRef: + properties: + apiVersion: + type: string + fieldPath: + type: string + required: + - fieldPath + type: object + x-kubernetes-map-type: atomic + fileKeyRef: + properties: + key: + type: string + optional: + default: false + type: boolean + path: + type: string + volumeName: + type: string + required: + - key + - path + - volumeName + type: object + x-kubernetes-map-type: atomic + resourceFieldRef: + properties: + containerName: + type: string + divisor: + anyOf: + - type: integer + - type: string + pattern: ^(\+|-)?(([0-9]+(\.[0-9]*)?)|(\.[0-9]+))(([KMGTPE]i)|[numkMGTPE]|([eE](\+|-)?(([0-9]+(\.[0-9]*)?)|(\.[0-9]+))))?$ + x-kubernetes-int-or-string: true + resource: + type: string + required: + - resource + type: object + x-kubernetes-map-type: atomic + secretKeyRef: + properties: + key: + type: string + name: + default: "" + type: string + optional: + type: boolean + required: + - key + type: object + x-kubernetes-map-type: atomic + type: object + required: + - name + type: object + maxItems: 128 + type: array + resources: + properties: + claims: + items: + properties: + name: + type: string + request: + type: string + required: + - name + type: object + type: array + x-kubernetes-list-map-keys: + - name + x-kubernetes-list-type: map + limits: + additionalProperties: + anyOf: + - type: integer + - type: string + pattern: ^(\+|-)?(([0-9]+(\.[0-9]*)?)|(\.[0-9]+))(([KMGTPE]i)|[numkMGTPE]|([eE](\+|-)?(([0-9]+(\.[0-9]*)?)|(\.[0-9]+))))?$ + x-kubernetes-int-or-string: true + type: object + requests: + additionalProperties: + anyOf: + - type: integer + - type: string + pattern: ^(\+|-)?(([0-9]+(\.[0-9]*)?)|(\.[0-9]+))(([KMGTPE]i)|[numkMGTPE]|([eE](\+|-)?(([0-9]+(\.[0-9]*)?)|(\.[0-9]+))))?$ + x-kubernetes-int-or-string: true + type: object + type: object + timeoutSeconds: + format: int64 + maximum: 3600 + minimum: 1 + type: integer + type: object + maxProperties: 16 + minProperties: 1 + type: object + required: + - actions + type: object volumeClaimTemplates: items: properties: diff --git a/config/rbac/role.yaml b/config/rbac/role.yaml index 22dac7bc..21e7862f 100644 --- a/config/rbac/role.yaml +++ b/config/rbac/role.yaml @@ -95,11 +95,20 @@ rules: - apps.wandb.com resources: - applications/status + - triageruns/status - weightsandbiases/status verbs: - get - patch - update +- apiGroups: + - apps.wandb.com + resources: + - triageruns + verbs: + - get + - list + - watch - apiGroups: - autoscaling resources: diff --git a/config/samples/apps_v2_application.yaml b/config/samples/apps_v2_application.yaml index 71714c26..181ebc5f 100644 --- a/config/samples/apps_v2_application.yaml +++ b/config/samples/apps_v2_application.yaml @@ -16,6 +16,20 @@ spec: - name: my-app image: nginx:1.14.2 + triage: + actions: + default: + containerName: my-app + command: ["nginx", "-t"] + timeoutSeconds: 60 + resources: + requests: + cpu: 25m + memory: 32Mi + limits: + cpu: 100m + memory: 128Mi + jobs: - metadata: name: my-job @@ -42,4 +56,3 @@ spec: - name: my-cronjob image: nginx:1.14.2 command: ["sleep", "10"] - diff --git a/deploy/operator/templates/wandb-operator-wandb-role.yaml b/deploy/operator/templates/wandb-operator-wandb-role.yaml index 290ae4e8..75a26d83 100644 --- a/deploy/operator/templates/wandb-operator-wandb-role.yaml +++ b/deploy/operator/templates/wandb-operator-wandb-role.yaml @@ -17,6 +17,14 @@ rules: - patch - update - watch + - apiGroups: + - apps.wandb.com + resources: + - triageruns + verbs: + - get + - list + - watch - apiGroups: - apps.wandb.com resources: @@ -28,6 +36,7 @@ rules: - apps.wandb.com resources: - applications/status + - triageruns/status - weightsandbiases/status verbs: - get @@ -210,4 +219,4 @@ subjects: - kind: ServiceAccount name: {{ include "wandb-operator.fullname" . }} namespace: {{ .Release.Namespace }} -{{- end }} \ No newline at end of file +{{- end }} diff --git a/internal/controller/reconciler/reconcile_v2.go b/internal/controller/reconciler/reconcile_v2.go index 57891d89..f9092eed 100644 --- a/internal/controller/reconciler/reconcile_v2.go +++ b/internal/controller/reconciler/reconcile_v2.go @@ -668,6 +668,7 @@ func reconcileApplications( setCustomCACertsChecksumAnnotation(&application.Spec.PodTemplate, caChecksum) application.Spec.HpaTemplate = ResolveAutoscaling(app, wandb) + application.Spec.Triage = resolveApplicationTriage(app.Triage) // Set shared service account for all W&B applications application.Spec.PodTemplate.Spec.ServiceAccountName = serviceAccountName diff --git a/internal/controller/reconciler/triage.go b/internal/controller/reconciler/triage.go new file mode 100644 index 00000000..1bb16721 --- /dev/null +++ b/internal/controller/reconciler/triage.go @@ -0,0 +1,34 @@ +package reconciler + +import ( + apiv2 "github.com/wandb/operator/api/v2" + serverManifest "github.com/wandb/operator/pkg/wandb/manifest" + corev1 "k8s.io/api/core/v1" +) + +func resolveApplicationTriage(triage *serverManifest.ApplicationTriage) *apiv2.ApplicationTriageSpec { + if triage == nil { + return nil + } + + actions := make(map[string]apiv2.TriageActionSpec, len(triage.Actions)) + for name, action := range triage.Actions { + env := make([]corev1.EnvVar, len(action.Env)) + for i := range action.Env { + env[i] = *action.Env[i].DeepCopy() + } + resolved := apiv2.TriageActionSpec{ + ContainerName: action.ContainerName, + Command: append([]string(nil), action.Command...), + Args: append([]string(nil), action.Args...), + Env: env, + TimeoutSeconds: action.TimeoutSeconds, + } + if action.Resources != nil { + resolved.Resources = action.Resources.DeepCopy() + } + actions[name] = resolved + } + + return &apiv2.ApplicationTriageSpec{Actions: actions} +} diff --git a/internal/controller/reconciler/triage_test.go b/internal/controller/reconciler/triage_test.go new file mode 100644 index 00000000..93f8c682 --- /dev/null +++ b/internal/controller/reconciler/triage_test.go @@ -0,0 +1,53 @@ +package reconciler + +import ( + "testing" + + serverManifest "github.com/wandb/operator/pkg/wandb/manifest" + corev1 "k8s.io/api/core/v1" + "k8s.io/apimachinery/pkg/api/resource" +) + +func TestResolveApplicationTriageCopiesCompactAction(t *testing.T) { + resources := &corev1.ResourceRequirements{ + Requests: corev1.ResourceList{ + corev1.ResourceCPU: resource.MustParse("100m"), + corev1.ResourceMemory: resource.MustParse("128Mi"), + }, + } + input := &serverManifest.ApplicationTriage{ + Actions: map[string]serverManifest.TriageAction{ + "default": { + ContainerName: "weave-trace", + Args: []string{"python", "-m", "weave_triage", "run-all", "--stream"}, + Env: []corev1.EnvVar{{ + Name: "PYTHONPATH", + Value: "/weave/src", + }}, + Resources: resources, + TimeoutSeconds: 600, + }, + }, + } + + resolved := resolveApplicationTriage(input) + action := resolved.Actions["default"] + if action.ContainerName != "weave-trace" { + t.Fatalf("containerName = %q", action.ContainerName) + } + if len(action.Args) != 5 || action.Args[4] != "--stream" { + t.Fatalf("args = %#v", action.Args) + } + if action.Resources == nil || action.Resources.Requests.Memory().String() != "128Mi" { + t.Fatalf("resources = %#v", action.Resources) + } + if action.TimeoutSeconds != 600 { + t.Fatalf("timeoutSeconds = %d", action.TimeoutSeconds) + } + + input.Actions["default"].Args[0] = "mutated" + resources.Requests[corev1.ResourceMemory] = resource.MustParse("1Gi") + if action.Args[0] != "python" || action.Resources.Requests.Memory().String() != "128Mi" { + t.Fatal("resolved action aliases mutable manifest data") + } +} diff --git a/internal/controller/triagerun_controller.go b/internal/controller/triagerun_controller.go new file mode 100644 index 00000000..37fc48c2 --- /dev/null +++ b/internal/controller/triagerun_controller.go @@ -0,0 +1,724 @@ +/* +Copyright 2025. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package controller + +import ( + "bufio" + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "strings" + "time" + + wandbv2 "github.com/wandb/operator/api/v2" + "github.com/wandb/operator/internal/controller/common" + batchv1 "k8s.io/api/batch/v1" + corev1 "k8s.io/api/core/v1" + apiextensionsv1 "k8s.io/apiextensions-apiserver/pkg/apis/apiextensions/v1" + apiequality "k8s.io/apimachinery/pkg/api/equality" + apierrors "k8s.io/apimachinery/pkg/api/errors" + apimeta "k8s.io/apimachinery/pkg/api/meta" + "k8s.io/apimachinery/pkg/api/resource" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/runtime" + "k8s.io/apimachinery/pkg/types" + corev1client "k8s.io/client-go/kubernetes/typed/core/v1" + ctrl "sigs.k8s.io/controller-runtime" + "sigs.k8s.io/controller-runtime/pkg/client" + "sigs.k8s.io/controller-runtime/pkg/controller/controllerutil" +) + +const ( + defaultTriageAction = "default" + defaultTriageTimeoutSeconds = int64(300) + maxTriageOutputBytes = int64(512 * 1024) + triageContainerName = "triage" + triageConditionSucceeded = "Succeeded" + triageRunLabel = "apps.wandb.com/triage-run" + triageApplicationLabel = "apps.wandb.com/triage-application" +) + +// TriagePodLogReader reads the structured output from a completed triage pod. +// It is an interface so controller behavior can be tested without a live API +// server's pod log subresource. +type TriagePodLogReader interface { + ReadPodLogs( + ctx context.Context, + namespace string, + podName string, + containerName string, + maxBytes int64, + ) ([]byte, error) +} + +type KubernetesTriagePodLogReader struct { + CoreV1 corev1client.CoreV1Interface +} + +func (r *KubernetesTriagePodLogReader) ReadPodLogs( + ctx context.Context, + namespace string, + podName string, + containerName string, + maxBytes int64, +) ([]byte, error) { + stream, err := r.CoreV1.Pods(namespace).GetLogs(podName, &corev1.PodLogOptions{ + Container: containerName, + }).Stream(ctx) + if err != nil { + return nil, err + } + defer stream.Close() + + output, err := io.ReadAll(io.LimitReader(stream, maxBytes+1)) + if err != nil { + return nil, err + } + if int64(len(output)) > maxBytes { + return nil, &triageOutputTooLargeError{maxBytes: maxBytes} + } + return output, nil +} + +// TriageRunReconciler turns each immutable TriageRun into exactly one Job and +// records the Job's structured JSONL output on the run status. +type TriageRunReconciler struct { + client.Client + Scheme *runtime.Scheme + PodLogs TriagePodLogReader +} + +// +kubebuilder:rbac:groups=apps.wandb.com,resources=triageruns,verbs=get;list;watch +// +kubebuilder:rbac:groups=apps.wandb.com,resources=triageruns/status,verbs=get;update;patch +// +kubebuilder:rbac:groups=apps.wandb.com,resources=applications,verbs=get;list;watch +// +kubebuilder:rbac:groups=batch,resources=jobs,verbs=get;list;watch;create +// +kubebuilder:rbac:groups="",resources=pods,verbs=get;list;watch +// +kubebuilder:rbac:groups="",resources=pods/log,verbs=get + +func (r *TriageRunReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) { + var run wandbv2.TriageRun + if err := r.Get(ctx, req.NamespacedName, &run); err != nil { + return ctrl.Result{}, client.IgnoreNotFound(err) + } + + if isTerminalTriagePhase(run.Status.Phase) { + return ctrl.Result{}, nil + } + + var application wandbv2.Application + applicationKey := types.NamespacedName{ + Namespace: run.Namespace, + Name: run.Spec.ApplicationRef.Name, + } + if err := r.Get(ctx, applicationKey, &application); err != nil { + if apierrors.IsNotFound(err) { + return r.failRun(ctx, &run, "ApplicationNotFound", + fmt.Sprintf("Application %q does not exist in namespace %q", applicationKey.Name, applicationKey.Namespace)) + } + return ctrl.Result{}, err + } + + actionName := run.Spec.Action + if actionName == "" { + actionName = defaultTriageAction + } + action, err := resolveTriageAction(&application, actionName) + if err != nil { + return r.failRun(ctx, &run, "InvalidAction", err.Error()) + } + + sourceContainer, err := selectTriageContainer(&application, action.ContainerName) + if err != nil { + return r.failRun(ctx, &run, "InvalidContainer", err.Error()) + } + + timeoutSeconds := action.TimeoutSeconds + if timeoutSeconds == 0 { + timeoutSeconds = defaultTriageTimeoutSeconds + } + resolved := resolvedTriageExecution(&application, sourceContainer, action, timeoutSeconds) + jobName := common.FitDefaultInfraName(run.Name, "-triage", 63) + + var job batchv1.Job + err = r.Get(ctx, types.NamespacedName{Namespace: run.Namespace, Name: jobName}, &job) + if err != nil && !apierrors.IsNotFound(err) { + return ctrl.Result{}, err + } + if apierrors.IsNotFound(err) { + if run.Status.JobRef != nil { + return r.failRun(ctx, &run, "JobMissing", + fmt.Sprintf("Job %q disappeared before the run completed", run.Status.JobRef.Name)) + } + + job = *buildTriageJob(&run, &application, sourceContainer, action, timeoutSeconds, jobName) + if err := controllerutil.SetControllerReference(&run, &job, r.Scheme); err != nil { + return ctrl.Result{}, err + } + if err := r.Create(ctx, &job); err != nil { + if !apierrors.IsAlreadyExists(err) { + return ctrl.Result{}, err + } + if err := r.Get(ctx, types.NamespacedName{Namespace: run.Namespace, Name: jobName}, &job); err != nil { + return ctrl.Result{}, err + } + } + } + + if !metav1.IsControlledBy(&job, &run) { + return r.failRun(ctx, &run, "JobNameCollision", + fmt.Sprintf("Job %q already exists and is not owned by this TriageRun", job.Name)) + } + + if jobFailed(&job) { + message := jobConditionMessage(&job, batchv1.JobFailed) + if message == "" { + message = fmt.Sprintf("Job %q failed", job.Name) + } + if results, collectErr := r.collectTriageResults(ctx, &job); collectErr == nil { + run.Status.Results = results + run.Status.Summary = summarizeTriageResults(results) + } + return r.failRunWithJob(ctx, &run, &job, resolved, "JobFailed", message) + } + + if jobComplete(&job) { + results, err := r.collectTriageResults(ctx, &job) + if err != nil { + var unavailable *triageOutputUnavailableError + if errors.As(err, &unavailable) { + return ctrl.Result{RequeueAfter: 2 * time.Second}, nil + } + return r.failRunWithJob(ctx, &run, &job, resolved, "InvalidOutput", err.Error()) + } + return r.completeRun(ctx, &run, &job, resolved, results) + } + + return r.markRunRunning(ctx, &run, &job, resolved) +} + +func resolveTriageAction(application *wandbv2.Application, actionName string) (wandbv2.TriageActionSpec, error) { + if application.Spec.Triage == nil { + return wandbv2.TriageActionSpec{}, fmt.Errorf("Application %q does not declare triage actions", application.Name) + } + action, ok := application.Spec.Triage.Actions[actionName] + if !ok { + return wandbv2.TriageActionSpec{}, fmt.Errorf( + "Application %q does not declare triage action %q", application.Name, actionName) + } + if len(action.Command) == 0 && len(action.Args) == 0 { + return wandbv2.TriageActionSpec{}, fmt.Errorf( + "triage action %q must override command or args", actionName) + } + return action, nil +} + +func selectTriageContainer(application *wandbv2.Application, name string) (*corev1.Container, error) { + containers := application.Spec.PodTemplate.Spec.Containers + if name == "" { + if len(containers) != 1 { + return nil, fmt.Errorf( + "triage action must select a container because Application %q has %d containers", + application.Name, len(containers)) + } + return &containers[0], nil + } + for i := range containers { + if containers[i].Name == name { + return &containers[i], nil + } + } + return nil, fmt.Errorf("Application %q has no container named %q", application.Name, name) +} + +func buildTriageJob( + run *wandbv2.TriageRun, + application *wandbv2.Application, + source *corev1.Container, + action wandbv2.TriageActionSpec, + timeoutSeconds int64, + jobName string, +) *batchv1.Job { + backoffLimit := int32(0) + podSpec := application.Spec.PodTemplate.Spec.DeepCopy() + podSpec.RestartPolicy = corev1.RestartPolicyNever + podSpec.InitContainers = nil + podSpec.EphemeralContainers = nil + podSpec.ReadinessGates = nil + podSpec.Containers = []corev1.Container{buildTriageContainer(source, action)} + + labels := map[string]string{ + triageRunLabel: common.FitDefaultInfraName(run.Name, "", 63), + triageApplicationLabel: common.FitDefaultInfraName(application.Name, "", 63), + } + return &batchv1.Job{ + ObjectMeta: metav1.ObjectMeta{ + Name: jobName, + Namespace: run.Namespace, + Labels: labels, + }, + Spec: batchv1.JobSpec{ + BackoffLimit: &backoffLimit, + ActiveDeadlineSeconds: &timeoutSeconds, + Template: corev1.PodTemplateSpec{ + ObjectMeta: metav1.ObjectMeta{Labels: labels}, + Spec: *podSpec, + }, + }, + } +} + +func buildTriageContainer(source *corev1.Container, action wandbv2.TriageActionSpec) corev1.Container { + container := corev1.Container{ + Name: triageContainerName, + Image: source.Image, + ImagePullPolicy: source.ImagePullPolicy, + Command: append([]string(nil), source.Command...), + Args: append([]string(nil), source.Args...), + WorkingDir: source.WorkingDir, + EnvFrom: append([]corev1.EnvFromSource(nil), source.EnvFrom...), + Env: mergeTriageEnv(source.Env, action.Env), + Resources: defaultTriageResources(), + VolumeMounts: append([]corev1.VolumeMount(nil), source.VolumeMounts...), + VolumeDevices: append([]corev1.VolumeDevice(nil), source.VolumeDevices...), + SecurityContext: source.SecurityContext.DeepCopy(), + TerminationMessagePath: source.TerminationMessagePath, + TerminationMessagePolicy: source.TerminationMessagePolicy, + } + if len(action.Command) > 0 { + container.Command = append([]string(nil), action.Command...) + } + if len(action.Args) > 0 { + container.Args = append([]string(nil), action.Args...) + } + if action.Resources != nil { + container.Resources = *action.Resources.DeepCopy() + } + return container +} + +func mergeTriageEnv(inherited, overrides []corev1.EnvVar) []corev1.EnvVar { + result := make([]corev1.EnvVar, 0, len(inherited)+len(overrides)) + overrideNames := make(map[string]struct{}, len(overrides)) + for _, env := range overrides { + overrideNames[env.Name] = struct{}{} + } + for _, env := range inherited { + if _, overridden := overrideNames[env.Name]; !overridden { + result = append(result, *env.DeepCopy()) + } + } + for _, env := range overrides { + result = append(result, *env.DeepCopy()) + } + return result +} + +func defaultTriageResources() corev1.ResourceRequirements { + return corev1.ResourceRequirements{ + Requests: corev1.ResourceList{ + corev1.ResourceCPU: resource.MustParse("100m"), + corev1.ResourceMemory: resource.MustParse("128Mi"), + }, + Limits: corev1.ResourceList{ + corev1.ResourceCPU: resource.MustParse("500m"), + corev1.ResourceMemory: resource.MustParse("512Mi"), + }, + } +} + +func resolvedTriageExecution( + application *wandbv2.Application, + source *corev1.Container, + action wandbv2.TriageActionSpec, + timeoutSeconds int64, +) *wandbv2.TriageResolvedExecution { + container := buildTriageContainer(source, action) + return &wandbv2.TriageResolvedExecution{ + ApplicationGeneration: application.Generation, + ContainerName: source.Name, + Image: container.Image, + Command: append([]string(nil), container.Command...), + Args: append([]string(nil), container.Args...), + TimeoutSeconds: timeoutSeconds, + } +} + +func (r *TriageRunReconciler) markRunRunning( + ctx context.Context, + run *wandbv2.TriageRun, + job *batchv1.Job, + resolved *wandbv2.TriageResolvedExecution, +) (ctrl.Result, error) { + statusBefore := run.DeepCopy().Status + run.Status.Phase = wandbv2.TriageRunPhaseRunning + run.Status.ObservedGeneration = run.Generation + run.Status.JobRef = &corev1.LocalObjectReference{Name: job.Name} + run.Status.ResolvedExecution = resolved + if run.Status.StartedAt == nil { + switch { + case job.Status.StartTime != nil: + run.Status.StartedAt = job.Status.StartTime.DeepCopy() + case !job.CreationTimestamp.IsZero(): + run.Status.StartedAt = job.CreationTimestamp.DeepCopy() + default: + now := metav1.Now() + run.Status.StartedAt = &now + } + } + apimeta.SetStatusCondition(&run.Status.Conditions, metav1.Condition{ + Type: triageConditionSucceeded, + Status: metav1.ConditionUnknown, + ObservedGeneration: run.Generation, + Reason: "JobRunning", + Message: fmt.Sprintf("Job %q is running", job.Name), + }) + if apiequality.Semantic.DeepEqual(statusBefore, run.Status) { + return ctrl.Result{}, nil + } + if err := r.Status().Update(ctx, run); err != nil { + return ctrl.Result{}, err + } + return ctrl.Result{}, nil +} + +func (r *TriageRunReconciler) completeRun( + ctx context.Context, + run *wandbv2.TriageRun, + job *batchv1.Job, + resolved *wandbv2.TriageResolvedExecution, + results []wandbv2.TriageCheckResult, +) (ctrl.Result, error) { + run.Status.Phase = wandbv2.TriageRunPhaseSucceeded + run.Status.ObservedGeneration = run.Generation + run.Status.JobRef = &corev1.LocalObjectReference{Name: job.Name} + run.Status.ResolvedExecution = resolved + run.Status.Results = results + run.Status.Summary = summarizeTriageResults(results) + run.Status.StartedAt = triageStartTime(run, job) + run.Status.CompletedAt = triageCompletionTime(job) + apimeta.SetStatusCondition(&run.Status.Conditions, metav1.Condition{ + Type: triageConditionSucceeded, + Status: metav1.ConditionTrue, + ObservedGeneration: run.Generation, + Reason: "ResultsCollected", + Message: fmt.Sprintf("Collected %d triage check results", len(results)), + }) + if err := r.Status().Update(ctx, run); err != nil { + return ctrl.Result{}, err + } + return ctrl.Result{}, nil +} + +func (r *TriageRunReconciler) failRun( + ctx context.Context, + run *wandbv2.TriageRun, + reason string, + message string, +) (ctrl.Result, error) { + return r.failRunWithJob(ctx, run, nil, run.Status.ResolvedExecution, reason, message) +} + +func (r *TriageRunReconciler) failRunWithJob( + ctx context.Context, + run *wandbv2.TriageRun, + job *batchv1.Job, + resolved *wandbv2.TriageResolvedExecution, + reason string, + message string, +) (ctrl.Result, error) { + run.Status.Phase = wandbv2.TriageRunPhaseFailed + run.Status.ObservedGeneration = run.Generation + run.Status.ResolvedExecution = resolved + now := metav1.Now() + run.Status.CompletedAt = &now + if job != nil { + run.Status.JobRef = &corev1.LocalObjectReference{Name: job.Name} + run.Status.StartedAt = triageStartTime(run, job) + if completedAt := triageCompletionTime(job); completedAt != nil { + run.Status.CompletedAt = completedAt + } + } + apimeta.SetStatusCondition(&run.Status.Conditions, metav1.Condition{ + Type: triageConditionSucceeded, + Status: metav1.ConditionFalse, + ObservedGeneration: run.Generation, + Reason: reason, + Message: message, + }) + if err := r.Status().Update(ctx, run); err != nil { + return ctrl.Result{}, err + } + return ctrl.Result{}, nil +} + +func triageStartTime(run *wandbv2.TriageRun, job *batchv1.Job) *metav1.Time { + if job.Status.StartTime != nil { + return job.Status.StartTime.DeepCopy() + } + if run.Status.StartedAt != nil { + return run.Status.StartedAt.DeepCopy() + } + if !job.CreationTimestamp.IsZero() { + return job.CreationTimestamp.DeepCopy() + } + return nil +} + +func triageCompletionTime(job *batchv1.Job) *metav1.Time { + if job.Status.CompletionTime != nil { + return job.Status.CompletionTime.DeepCopy() + } + now := metav1.Now() + return &now +} + +func jobComplete(job *batchv1.Job) bool { + return jobConditionTrue(job, batchv1.JobComplete) +} + +func jobFailed(job *batchv1.Job) bool { + return jobConditionTrue(job, batchv1.JobFailed) +} + +func jobConditionTrue(job *batchv1.Job, conditionType batchv1.JobConditionType) bool { + for _, condition := range job.Status.Conditions { + if condition.Type == conditionType && condition.Status == corev1.ConditionTrue { + return true + } + } + return false +} + +func jobConditionMessage(job *batchv1.Job, conditionType batchv1.JobConditionType) string { + for _, condition := range job.Status.Conditions { + if condition.Type == conditionType && condition.Status == corev1.ConditionTrue { + return condition.Message + } + } + return "" +} + +func isTerminalTriagePhase(phase wandbv2.TriageRunPhase) bool { + return phase == wandbv2.TriageRunPhaseSucceeded || phase == wandbv2.TriageRunPhaseFailed +} + +func (r *TriageRunReconciler) collectTriageResults( + ctx context.Context, + job *batchv1.Job, +) ([]wandbv2.TriageCheckResult, error) { + if r.PodLogs == nil { + return nil, errors.New("pod log reader is not configured") + } + + var pods corev1.PodList + if err := r.List(ctx, &pods, + client.InNamespace(job.Namespace), + client.MatchingLabels{"batch.kubernetes.io/job-name": job.Name}, + ); err != nil { + return nil, &triageOutputUnavailableError{ + err: fmt.Errorf("list pods for Job %q: %w", job.Name, err), + } + } + if len(pods.Items) == 0 { + return nil, &triageOutputUnavailableError{ + err: fmt.Errorf("pod for Job %q is not available yet", job.Name), + } + } + if len(pods.Items) > 1 { + return nil, fmt.Errorf("expected one pod for Job %q, found %d", job.Name, len(pods.Items)) + } + + output, err := r.PodLogs.ReadPodLogs( + ctx, job.Namespace, pods.Items[0].Name, triageContainerName, maxTriageOutputBytes) + if err != nil { + var tooLarge *triageOutputTooLargeError + if errors.As(err, &tooLarge) { + return nil, err + } + return nil, &triageOutputUnavailableError{ + err: fmt.Errorf("read triage output: %w", err), + } + } + results, err := parseTriageJSONL(output) + if err != nil { + return nil, fmt.Errorf("parse triage output: %w", err) + } + return results, nil +} + +type triageOutputUnavailableError struct { + err error +} + +func (e *triageOutputUnavailableError) Error() string { + return e.err.Error() +} + +func (e *triageOutputUnavailableError) Unwrap() error { + return e.err +} + +type triageOutputTooLargeError struct { + maxBytes int64 +} + +func (e *triageOutputTooLargeError) Error() string { + return fmt.Sprintf("triage output exceeds %d bytes", e.maxBytes) +} + +type triageJSONResult struct { + Name string `json:"name"` + Umbrella string `json:"umbrella,omitempty"` + Severity string `json:"severity"` + Message string `json:"message,omitempty"` + Evidence json.RawMessage `json:"evidence,omitempty"` + Remediation string `json:"remediation,omitempty"` + StartedAt string `json:"started_at,omitempty"` + EndedAt string `json:"ended_at,omitempty"` + DurationMS int64 `json:"duration_ms,omitempty"` +} + +func parseTriageJSONL(output []byte) ([]wandbv2.TriageCheckResult, error) { + if int64(len(output)) > maxTriageOutputBytes { + return nil, fmt.Errorf("triage output exceeds %d bytes", maxTriageOutputBytes) + } + + scanner := bufio.NewScanner(bytes.NewReader(output)) + scanner.Buffer(make([]byte, 64*1024), int(maxTriageOutputBytes)) + results := make([]wandbv2.TriageCheckResult, 0) + lineNumber := 0 + for scanner.Scan() { + lineNumber++ + line := bytes.TrimSpace(scanner.Bytes()) + if len(line) == 0 { + continue + } + + var raw triageJSONResult + decoder := json.NewDecoder(bytes.NewReader(line)) + if err := decoder.Decode(&raw); err != nil { + return nil, fmt.Errorf("line %d is not valid JSON: %w", lineNumber, err) + } + var trailing any + if err := decoder.Decode(&trailing); !errors.Is(err, io.EOF) { + return nil, fmt.Errorf("line %d contains more than one JSON value", lineNumber) + } + result, err := raw.toAPIResult() + if err != nil { + return nil, fmt.Errorf("line %d: %w", lineNumber, err) + } + results = append(results, result) + } + if err := scanner.Err(); err != nil { + return nil, err + } + if len(results) == 0 { + return nil, errors.New("triage command emitted no check results") + } + return results, nil +} + +func (r triageJSONResult) toAPIResult() (wandbv2.TriageCheckResult, error) { + if strings.TrimSpace(r.Name) == "" { + return wandbv2.TriageCheckResult{}, errors.New("name is required") + } + severity := wandbv2.TriageSeverity(r.Severity) + switch severity { + case wandbv2.TriageSeverityPass, + wandbv2.TriageSeverityWarn, + wandbv2.TriageSeverityFail, + wandbv2.TriageSeverityError: + default: + return wandbv2.TriageCheckResult{}, fmt.Errorf("unsupported severity %q", r.Severity) + } + + result := wandbv2.TriageCheckResult{ + Name: r.Name, + Umbrella: r.Umbrella, + Severity: severity, + Message: r.Message, + Remediation: r.Remediation, + DurationMilliseconds: r.DurationMS, + } + if len(r.Evidence) > 0 && !bytes.Equal(r.Evidence, []byte("null")) { + if !json.Valid(r.Evidence) { + return wandbv2.TriageCheckResult{}, errors.New("evidence is not valid JSON") + } + result.Evidence = &apiextensionsv1.JSON{Raw: append([]byte(nil), r.Evidence...)} + } + + var err error + result.StartedAt, err = parseTriageTimestamp("started_at", r.StartedAt) + if err != nil { + return wandbv2.TriageCheckResult{}, err + } + result.EndedAt, err = parseTriageTimestamp("ended_at", r.EndedAt) + if err != nil { + return wandbv2.TriageCheckResult{}, err + } + return result, nil +} + +func parseTriageTimestamp(field, value string) (*metav1.Time, error) { + if value == "" { + return nil, nil + } + parsed, err := time.Parse(time.RFC3339Nano, value) + if err != nil { + return nil, fmt.Errorf("%s must be RFC3339: %w", field, err) + } + timestamp := metav1.NewTime(parsed) + return ×tamp, nil +} + +func summarizeTriageResults(results []wandbv2.TriageCheckResult) *wandbv2.TriageRunSummary { + summary := &wandbv2.TriageRunSummary{Total: int32(len(results))} + for _, result := range results { + switch result.Severity { + case wandbv2.TriageSeverityPass: + summary.Pass++ + case wandbv2.TriageSeverityWarn: + summary.Warn++ + case wandbv2.TriageSeverityFail: + summary.Fail++ + case wandbv2.TriageSeverityError: + summary.Error++ + } + } + switch { + case summary.Error > 0: + summary.OverallSeverity = wandbv2.TriageSeverityError + case summary.Fail > 0: + summary.OverallSeverity = wandbv2.TriageSeverityFail + case summary.Warn > 0: + summary.OverallSeverity = wandbv2.TriageSeverityWarn + default: + summary.OverallSeverity = wandbv2.TriageSeverityPass + } + return summary +} + +func (r *TriageRunReconciler) SetupWithManager(mgr ctrl.Manager) error { + return ctrl.NewControllerManagedBy(mgr). + For(&wandbv2.TriageRun{}). + Owns(&batchv1.Job{}). + Named("triagerun"). + Complete(r) +} diff --git a/internal/controller/triagerun_controller_unit_test.go b/internal/controller/triagerun_controller_unit_test.go new file mode 100644 index 00000000..0cb094d5 --- /dev/null +++ b/internal/controller/triagerun_controller_unit_test.go @@ -0,0 +1,340 @@ +package controller + +import ( + "context" + "fmt" + "testing" + + wandbv2 "github.com/wandb/operator/api/v2" + serverManifest "github.com/wandb/operator/pkg/wandb/manifest" + batchv1 "k8s.io/api/batch/v1" + corev1 "k8s.io/api/core/v1" + "k8s.io/apimachinery/pkg/api/resource" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/runtime" + "k8s.io/apimachinery/pkg/types" + ctrl "sigs.k8s.io/controller-runtime" + "sigs.k8s.io/controller-runtime/pkg/client" + "sigs.k8s.io/controller-runtime/pkg/client/fake" + "sigs.k8s.io/controller-runtime/pkg/controller/controllerutil" + "sigs.k8s.io/yaml" +) + +func TestTriageRunCreatesBoundedJobFromApplication(t *testing.T) { + t.Parallel() + + testScheme := newTriageTestScheme(t) + run := testTriageRun() + application := testTriageApplication() + fakeClient := fake.NewClientBuilder(). + WithScheme(testScheme). + WithStatusSubresource(&wandbv2.TriageRun{}, &batchv1.Job{}). + WithObjects(run, application). + Build() + reconciler := &TriageRunReconciler{ + Client: fakeClient, + Scheme: testScheme, + PodLogs: &staticTriageLogReader{}, + } + + _, err := reconciler.Reconcile(context.Background(), requestFor(run)) + if err != nil { + t.Fatalf("reconcile TriageRun: %v", err) + } + + var job batchv1.Job + if err := fakeClient.Get(context.Background(), types.NamespacedName{ + Namespace: run.Namespace, + Name: "weave-check-triage", + }, &job); err != nil { + t.Fatalf("get triage Job: %v", err) + } + if job.Spec.Template.Spec.ServiceAccountName != application.Spec.PodTemplate.Spec.ServiceAccountName { + t.Fatalf("service account = %q, want parent application SA %q", + job.Spec.Template.Spec.ServiceAccountName, + application.Spec.PodTemplate.Spec.ServiceAccountName) + } + if job.Spec.BackoffLimit == nil || *job.Spec.BackoffLimit != 0 { + t.Fatalf("backoffLimit = %v, want 0", job.Spec.BackoffLimit) + } + if job.Spec.Template.Spec.RestartPolicy != corev1.RestartPolicyNever { + t.Fatalf("restartPolicy = %q, want Never", job.Spec.Template.Spec.RestartPolicy) + } + if len(job.Spec.Template.Spec.Containers) != 1 { + t.Fatalf("containers = %d, want exactly one", len(job.Spec.Template.Spec.Containers)) + } + container := job.Spec.Template.Spec.Containers[0] + if container.Image != "weave:sha256-test" { + t.Fatalf("image = %q, want inherited image", container.Image) + } + if got := container.Resources.Requests.Cpu().String(); got != "100m" { + t.Fatalf("CPU request = %q, want small default 100m", got) + } + if got := container.Resources.Requests.Memory().String(); got != "128Mi" { + t.Fatalf("memory request = %q, want small default 128Mi", got) + } + if got := application.Spec.PodTemplate.Spec.Containers[0].Resources.Requests.Memory().String(); got != "8Gi" { + t.Fatalf("test parent memory request = %q, want 8Gi", got) + } + if container.Resources.Requests.Memory().Cmp( + *application.Spec.PodTemplate.Spec.Containers[0].Resources.Requests.Memory(), + ) >= 0 { + t.Fatal("triage memory request must be smaller than the parent application request") + } + if len(container.Args) != 4 || container.Args[0] != "python" { + t.Fatalf("args = %#v, want triage action args", container.Args) + } + if got := envValue(container.Env, "PYTHONPATH"); got != "/weave/src" { + t.Fatalf("PYTHONPATH = %q, want action override", got) + } + if got := envValue(container.Env, "DATABASE_URL"); got != "mysql://wandb" { + t.Fatalf("DATABASE_URL = %q, want inherited env", got) + } + if len(container.Ports) != 0 || container.ReadinessProbe != nil || container.LivenessProbe != nil { + t.Fatal("triage container must not inherit serving ports or probes") + } + if len(container.VolumeMounts) != 1 || len(job.Spec.Template.Spec.Volumes) != 1 { + t.Fatal("triage Job must inherit the selected container's mounts and parent pod volumes") + } + if !metav1.IsControlledBy(&job, run) { + t.Fatal("triage Job is not controlled by the TriageRun") + } + + var updatedRun wandbv2.TriageRun + if err := fakeClient.Get(context.Background(), client.ObjectKeyFromObject(run), &updatedRun); err != nil { + t.Fatalf("get updated TriageRun: %v", err) + } + if updatedRun.Status.Phase != wandbv2.TriageRunPhaseRunning { + t.Fatalf("phase = %q, want Running", updatedRun.Status.Phase) + } + if updatedRun.Status.JobRef == nil || updatedRun.Status.JobRef.Name != job.Name { + t.Fatalf("jobRef = %#v, want %q", updatedRun.Status.JobRef, job.Name) + } + if updatedRun.Status.ResolvedExecution == nil || + updatedRun.Status.ResolvedExecution.ContainerName != "weave-trace" { + t.Fatalf("resolved execution = %#v, want weave-trace container", updatedRun.Status.ResolvedExecution) + } +} + +func TestTriageRunCollectsFailedCheckAsSuccessfulExecution(t *testing.T) { + t.Parallel() + + testScheme := newTriageTestScheme(t) + run := testTriageRun() + run.Status.Phase = wandbv2.TriageRunPhaseRunning + run.Status.JobRef = &corev1.LocalObjectReference{Name: "weave-check-triage"} + application := testTriageApplication() + action := application.Spec.Triage.Actions[defaultTriageAction] + source := &application.Spec.PodTemplate.Spec.Containers[0] + job := buildTriageJob( + run, application, source, action, defaultTriageTimeoutSeconds, "weave-check-triage") + if err := controllerutil.SetControllerReference(run, job, testScheme); err != nil { + t.Fatalf("set Job owner: %v", err) + } + job.Status.Conditions = []batchv1.JobCondition{{ + Type: batchv1.JobComplete, + Status: corev1.ConditionTrue, + }} + pod := &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Name: "weave-check-triage-abcde", + Namespace: run.Namespace, + Labels: map[string]string{"batch.kubernetes.io/job-name": job.Name}, + }, + } + logs := []byte( + `{"name":"starter-project","severity":"pass","message":"reachable"}` + "\n" + + `{"name":"starter-object","severity":"fail","evidence":{"missing":2},"remediation":"create starters"}` + "\n", + ) + fakeClient := fake.NewClientBuilder(). + WithScheme(testScheme). + WithStatusSubresource(&wandbv2.TriageRun{}, &batchv1.Job{}). + WithObjects(run, application, job, pod). + Build() + reconciler := &TriageRunReconciler{ + Client: fakeClient, + Scheme: testScheme, + PodLogs: &staticTriageLogReader{output: logs}, + } + + _, err := reconciler.Reconcile(context.Background(), requestFor(run)) + if err != nil { + t.Fatalf("reconcile completed TriageRun: %v", err) + } + + var updatedRun wandbv2.TriageRun + if err := fakeClient.Get(context.Background(), client.ObjectKeyFromObject(run), &updatedRun); err != nil { + t.Fatalf("get updated TriageRun: %v", err) + } + if updatedRun.Status.Phase != wandbv2.TriageRunPhaseSucceeded { + t.Fatalf("phase = %q, want Succeeded because the command completed", updatedRun.Status.Phase) + } + if updatedRun.Status.Summary == nil || + updatedRun.Status.Summary.OverallSeverity != wandbv2.TriageSeverityFail || + updatedRun.Status.Summary.Fail != 1 { + t.Fatalf("summary = %#v, want one failed diagnostic check", updatedRun.Status.Summary) + } + if len(updatedRun.Status.Results) != 2 { + t.Fatalf("results = %d, want 2", len(updatedRun.Status.Results)) + } +} + +func TestParseTriageJSONLRejectsNoisyOutput(t *testing.T) { + t.Parallel() + + _, err := parseTriageJSONL([]byte( + `{"name":"starter-project","severity":"pass"}` + "\n" + + "debug: checking object\n", + )) + if err == nil { + t.Fatal("expected non-JSON output to be rejected") + } +} + +func TestManifestTriageActionDecodes(t *testing.T) { + t.Parallel() + + var decoded serverManifest.Manifest + input := []byte(` +applications: + weave-trace: + triage: + actions: + default: + containerName: weave-trace + args: [python, -m, weave_triage, run-all, --stream] + timeoutSeconds: 600 + resources: + requests: + cpu: 100m + memory: 128Mi +`) + if err := yaml.Unmarshal(input, &decoded); err != nil { + t.Fatalf("decode manifest triage action: %v", err) + } + action := decoded.Applications["weave-trace"].Triage.Actions["default"] + if action.ContainerName != "weave-trace" || action.TimeoutSeconds != 600 { + t.Fatalf("decoded action = %#v", action) + } +} + +type staticTriageLogReader struct { + output []byte + err error +} + +func (r *staticTriageLogReader) ReadPodLogs( + _ context.Context, + _ string, + _ string, + _ string, + _ int64, +) ([]byte, error) { + return r.output, r.err +} + +func newTriageTestScheme(t *testing.T) *runtime.Scheme { + t.Helper() + testScheme := runtime.NewScheme() + for name, add := range map[string]func(*runtime.Scheme) error{ + "apps.wandb.com/v2": wandbv2.AddToScheme, + "batch/v1": batchv1.AddToScheme, + "core/v1": corev1.AddToScheme, + } { + if err := add(testScheme); err != nil { + t.Fatalf("add %s to scheme: %v", name, err) + } + } + return testScheme +} + +func testTriageRun() *wandbv2.TriageRun { + return &wandbv2.TriageRun{ + TypeMeta: metav1.TypeMeta{ + APIVersion: wandbv2.GroupVersion.String(), + Kind: "TriageRun", + }, + ObjectMeta: metav1.ObjectMeta{ + Name: "weave-check", + Namespace: "wandb", + UID: types.UID("triage-run-uid"), + }, + Spec: wandbv2.TriageRunSpec{ + ApplicationRef: wandbv2.TriageApplicationReference{Name: "weave-trace"}, + Action: defaultTriageAction, + }, + } +} + +func testTriageApplication() *wandbv2.Application { + return &wandbv2.Application{ + ObjectMeta: metav1.ObjectMeta{ + Name: "weave-trace", + Namespace: "wandb", + Generation: 7, + }, + Spec: wandbv2.ApplicationSpec{ + PodTemplate: corev1.PodTemplateSpec{ + Spec: corev1.PodSpec{ + ServiceAccountName: "wandb-app", + ImagePullSecrets: []corev1.LocalObjectReference{{Name: "registry"}}, + Volumes: []corev1.Volume{{ + Name: "ca", + VolumeSource: corev1.VolumeSource{ + ConfigMap: &corev1.ConfigMapVolumeSource{ + LocalObjectReference: corev1.LocalObjectReference{Name: "custom-ca"}, + }, + }, + }}, + Containers: []corev1.Container{{ + Name: "weave-trace", + Image: "weave:sha256-test", + Args: []string{"uvicorn", "weave.trace_server.app:app"}, + Env: []corev1.EnvVar{ + {Name: "DATABASE_URL", Value: "mysql://wandb"}, + {Name: "PYTHONPATH", Value: "/parent"}, + }, + Resources: corev1.ResourceRequirements{ + Requests: corev1.ResourceList{ + corev1.ResourceCPU: resource.MustParse("1"), + corev1.ResourceMemory: resource.MustParse("8Gi"), + }, + }, + Ports: []corev1.ContainerPort{{ContainerPort: 8080}}, + ReadinessProbe: &corev1.Probe{}, + LivenessProbe: &corev1.Probe{}, + VolumeMounts: []corev1.VolumeMount{{ + Name: "ca", + MountPath: "/etc/ssl/custom-ca.pem", + }}, + }}, + }, + }, + Triage: &wandbv2.ApplicationTriageSpec{ + Actions: map[string]wandbv2.TriageActionSpec{ + defaultTriageAction: { + ContainerName: "weave-trace", + Args: []string{"python", "-m", "weave_triage", "run-all"}, + Env: []corev1.EnvVar{{ + Name: "PYTHONPATH", + Value: "/weave/src", + }}, + }, + }, + }, + }, + } +} + +func requestFor(run *wandbv2.TriageRun) ctrl.Request { + return ctrl.Request{NamespacedName: client.ObjectKeyFromObject(run)} +} + +func envValue(env []corev1.EnvVar, name string) string { + for _, variable := range env { + if variable.Name == name { + return variable.Value + } + } + return fmt.Sprintf("<%s not found>", name) +} diff --git a/internal/crdinstaller/crds/operator/apps.wandb.com_applications.yaml b/internal/crdinstaller/crds/operator/apps.wandb.com_applications.yaml index 67be4bc9..5cd5c9dd 100644 --- a/internal/crdinstaller/crds/operator/apps.wandb.com_applications.yaml +++ b/internal/crdinstaller/crds/operator/apps.wandb.com_applications.yaml @@ -12801,6 +12801,154 @@ spec: type: type: string type: object + triage: + properties: + actions: + additionalProperties: + properties: + args: + items: + type: string + maxItems: 64 + minItems: 1 + type: array + command: + items: + type: string + maxItems: 64 + minItems: 1 + type: array + containerName: + type: string + env: + items: + properties: + name: + type: string + value: + type: string + valueFrom: + properties: + configMapKeyRef: + properties: + key: + type: string + name: + default: "" + type: string + optional: + type: boolean + required: + - key + type: object + x-kubernetes-map-type: atomic + fieldRef: + properties: + apiVersion: + type: string + fieldPath: + type: string + required: + - fieldPath + type: object + x-kubernetes-map-type: atomic + fileKeyRef: + properties: + key: + type: string + optional: + default: false + type: boolean + path: + type: string + volumeName: + type: string + required: + - key + - path + - volumeName + type: object + x-kubernetes-map-type: atomic + resourceFieldRef: + properties: + containerName: + type: string + divisor: + anyOf: + - type: integer + - type: string + pattern: ^(\+|-)?(([0-9]+(\.[0-9]*)?)|(\.[0-9]+))(([KMGTPE]i)|[numkMGTPE]|([eE](\+|-)?(([0-9]+(\.[0-9]*)?)|(\.[0-9]+))))?$ + x-kubernetes-int-or-string: true + resource: + type: string + required: + - resource + type: object + x-kubernetes-map-type: atomic + secretKeyRef: + properties: + key: + type: string + name: + default: "" + type: string + optional: + type: boolean + required: + - key + type: object + x-kubernetes-map-type: atomic + type: object + required: + - name + type: object + maxItems: 128 + type: array + resources: + properties: + claims: + items: + properties: + name: + type: string + request: + type: string + required: + - name + type: object + type: array + x-kubernetes-list-map-keys: + - name + x-kubernetes-list-type: map + limits: + additionalProperties: + anyOf: + - type: integer + - type: string + pattern: ^(\+|-)?(([0-9]+(\.[0-9]*)?)|(\.[0-9]+))(([KMGTPE]i)|[numkMGTPE]|([eE](\+|-)?(([0-9]+(\.[0-9]*)?)|(\.[0-9]+))))?$ + x-kubernetes-int-or-string: true + type: object + requests: + additionalProperties: + anyOf: + - type: integer + - type: string + pattern: ^(\+|-)?(([0-9]+(\.[0-9]*)?)|(\.[0-9]+))(([KMGTPE]i)|[numkMGTPE]|([eE](\+|-)?(([0-9]+(\.[0-9]*)?)|(\.[0-9]+))))?$ + x-kubernetes-int-or-string: true + type: object + type: object + timeoutSeconds: + format: int64 + maximum: 3600 + minimum: 1 + type: integer + type: object + maxProperties: 16 + minProperties: 1 + type: object + required: + - actions + type: object volumeClaimTemplates: items: properties: diff --git a/pkg/wandb/manifest/manifest.go b/pkg/wandb/manifest/manifest.go index 1cf06bd7..f8070905 100644 --- a/pkg/wandb/manifest/manifest.go +++ b/pkg/wandb/manifest/manifest.go @@ -175,6 +175,23 @@ type Application struct { VolumeMounts []VolumeMount `yaml:"volumeMounts,omitempty"` Sizing map[v2.Size]SizingConfig `yaml:"sizing,omitempty"` Ingress *AppIngressSpec `yaml:"ingress,omitempty"` + Triage *ApplicationTriage `yaml:"triage,omitempty"` +} + +// ApplicationTriage declares diagnostic actions available for an application. +// The operator copies these compact overrides onto the generated Application +// CR, which remains the source of the runtime pod configuration. +type ApplicationTriage struct { + Actions map[string]TriageAction `yaml:"actions"` +} + +type TriageAction struct { + ContainerName string `yaml:"containerName,omitempty"` + Command []string `yaml:"command,omitempty"` + Args []string `yaml:"args,omitempty"` + Env []corev1.EnvVar `yaml:"env,omitempty"` + Resources *corev1.ResourceRequirements `yaml:"resources,omitempty"` + TimeoutSeconds int64 `yaml:"timeoutSeconds,omitempty"` } type AppIngressSpec struct { From 8f4923da8bc411a0e385edc413722ff7fd6604d8 Mon Sep 17 00:00:00 2001 From: Aravind Warrier Date: Tue, 4 Aug 2026 11:43:00 -0400 Subject: [PATCH 2/4] feat(controller): Reconcile multiple triage actions --- internal/controller/triagerun_controller.go | 396 ++++++++++++------ .../triagerun_controller_unit_test.go | 133 +++++- 2 files changed, 390 insertions(+), 139 deletions(-) diff --git a/internal/controller/triagerun_controller.go b/internal/controller/triagerun_controller.go index 37fc48c2..7d1fed5a 100644 --- a/internal/controller/triagerun_controller.go +++ b/internal/controller/triagerun_controller.go @@ -24,6 +24,7 @@ import ( "errors" "fmt" "io" + "strconv" "strings" "time" @@ -53,6 +54,7 @@ const ( triageConditionSucceeded = "Succeeded" triageRunLabel = "apps.wandb.com/triage-run" triageApplicationLabel = "apps.wandb.com/triage-application" + triageActionAnnotation = "apps.wandb.com/triage-action" ) // TriagePodLogReader reads the structured output from a completed triage pod. @@ -97,8 +99,8 @@ func (r *KubernetesTriagePodLogReader) ReadPodLogs( return output, nil } -// TriageRunReconciler turns each immutable TriageRun into exactly one Job and -// records the Job's structured JSONL output on the run status. +// TriageRunReconciler turns each selected action on an immutable TriageRun into +// one Job and records every Job's structured JSONL output on the run status. type TriageRunReconciler struct { client.Client Scheme *runtime.Scheme @@ -135,82 +137,153 @@ func (r *TriageRunReconciler) Reconcile(ctx context.Context, req ctrl.Request) ( return ctrl.Result{}, err } - actionName := run.Spec.Action - if actionName == "" { - actionName = defaultTriageAction - } - action, err := resolveTriageAction(&application, actionName) + actions, err := resolveRequestedTriageActions(&run, &application) if err != nil { return r.failRun(ctx, &run, "InvalidAction", err.Error()) } - sourceContainer, err := selectTriageContainer(&application, action.ContainerName) - if err != nil { - return r.failRun(ctx, &run, "InvalidContainer", err.Error()) + statusBefore := run.DeepCopy().Status + statuses := make([]wandbv2.TriageActionStatus, 0, len(actions)) + requeueForOutput := false + failureMessages := make([]string, 0) + for i := range actions { + action := &actions[i] + previous := findTriageActionStatus(run.Status.ActionStatuses, action.name) + job, err := r.getOrCreateTriageJob(ctx, &run, &application, action, previous) + if err != nil { + var missing *triageJobMissingError + if errors.As(err, &missing) { + return r.failRun(ctx, &run, "JobMissing", missing.Error()) + } + return ctrl.Result{}, err + } + if !metav1.IsControlledBy(job, &run) { + return r.failRun(ctx, &run, "JobNameCollision", + fmt.Sprintf("Job %q already exists and is not owned by this TriageRun", job.Name)) + } + + status, outputUnavailable, failureMessage := r.actionStatus(ctx, action, job, previous) + statuses = append(statuses, status) + requeueForOutput = requeueForOutput || outputUnavailable + if failureMessage != "" { + failureMessages = append(failureMessages, failureMessage) + } } - timeoutSeconds := action.TimeoutSeconds - if timeoutSeconds == 0 { - timeoutSeconds = defaultTriageTimeoutSeconds + run.Status.ActionStatuses = statuses + setAggregateTriageStatus(&run, failureMessages) + if !apiequality.Semantic.DeepEqual(statusBefore, run.Status) { + if err := r.Status().Update(ctx, &run); err != nil { + return ctrl.Result{}, err + } } - resolved := resolvedTriageExecution(&application, sourceContainer, action, timeoutSeconds) - jobName := common.FitDefaultInfraName(run.Name, "-triage", 63) + if requeueForOutput { + return ctrl.Result{RequeueAfter: 2 * time.Second}, nil + } + return ctrl.Result{}, nil +} + +type resolvedTriageAction struct { + name wandbv2.TriageActionName + spec wandbv2.TriageActionSpec + source *corev1.Container + timeoutSeconds int64 + resolved *wandbv2.TriageResolvedExecution + jobName string +} + +func resolveRequestedTriageActions( + run *wandbv2.TriageRun, + application *wandbv2.Application, +) ([]resolvedTriageAction, error) { + if len(run.Spec.Actions) == 0 { + return nil, errors.New("at least one triage action is required") + } + actions := make([]resolvedTriageAction, 0, len(run.Spec.Actions)) + for i, actionName := range run.Spec.Actions { + spec, err := resolveTriageAction(application, string(actionName)) + if err != nil { + return nil, err + } + source, err := selectTriageContainer(application, spec.ContainerName) + if err != nil { + return nil, err + } + timeoutSeconds := spec.TimeoutSeconds + if timeoutSeconds == 0 { + timeoutSeconds = defaultTriageTimeoutSeconds + } + actions = append(actions, resolvedTriageAction{ + name: actionName, + spec: spec, + source: source, + timeoutSeconds: timeoutSeconds, + resolved: resolvedTriageExecution(application, source, spec, timeoutSeconds), + jobName: common.FitDefaultInfraName( + run.Name, "-triage-"+strconv.Itoa(i), 63), + }) + } + return actions, nil +} + +func findTriageActionStatus( + statuses []wandbv2.TriageActionStatus, + action wandbv2.TriageActionName, +) *wandbv2.TriageActionStatus { + for i := range statuses { + if statuses[i].Action == action { + return &statuses[i] + } + } + return nil +} +func (r *TriageRunReconciler) getOrCreateTriageJob( + ctx context.Context, + run *wandbv2.TriageRun, + application *wandbv2.Application, + action *resolvedTriageAction, + previous *wandbv2.TriageActionStatus, +) (*batchv1.Job, error) { var job batchv1.Job - err = r.Get(ctx, types.NamespacedName{Namespace: run.Namespace, Name: jobName}, &job) + err := r.Get(ctx, types.NamespacedName{Namespace: run.Namespace, Name: action.jobName}, &job) if err != nil && !apierrors.IsNotFound(err) { - return ctrl.Result{}, err + return nil, err } if apierrors.IsNotFound(err) { - if run.Status.JobRef != nil { - return r.failRun(ctx, &run, "JobMissing", - fmt.Sprintf("Job %q disappeared before the run completed", run.Status.JobRef.Name)) + if previous != nil && previous.JobRef != nil { + return nil, &triageJobMissingError{ + name: previous.JobRef.Name, action: action.name, + } } - - job = *buildTriageJob(&run, &application, sourceContainer, action, timeoutSeconds, jobName) - if err := controllerutil.SetControllerReference(&run, &job, r.Scheme); err != nil { - return ctrl.Result{}, err + job = *buildTriageJob( + run, application, action.source, action.spec, action.timeoutSeconds, action.jobName) + job.Annotations = map[string]string{triageActionAnnotation: string(action.name)} + if err := controllerutil.SetControllerReference(run, &job, r.Scheme); err != nil { + return nil, err } if err := r.Create(ctx, &job); err != nil { if !apierrors.IsAlreadyExists(err) { - return ctrl.Result{}, err + return nil, err } - if err := r.Get(ctx, types.NamespacedName{Namespace: run.Namespace, Name: jobName}, &job); err != nil { - return ctrl.Result{}, err + if err := r.Get(ctx, types.NamespacedName{ + Namespace: run.Namespace, + Name: action.jobName, + }, &job); err != nil { + return nil, err } } } + return &job, nil +} - if !metav1.IsControlledBy(&job, &run) { - return r.failRun(ctx, &run, "JobNameCollision", - fmt.Sprintf("Job %q already exists and is not owned by this TriageRun", job.Name)) - } - - if jobFailed(&job) { - message := jobConditionMessage(&job, batchv1.JobFailed) - if message == "" { - message = fmt.Sprintf("Job %q failed", job.Name) - } - if results, collectErr := r.collectTriageResults(ctx, &job); collectErr == nil { - run.Status.Results = results - run.Status.Summary = summarizeTriageResults(results) - } - return r.failRunWithJob(ctx, &run, &job, resolved, "JobFailed", message) - } - - if jobComplete(&job) { - results, err := r.collectTriageResults(ctx, &job) - if err != nil { - var unavailable *triageOutputUnavailableError - if errors.As(err, &unavailable) { - return ctrl.Result{RequeueAfter: 2 * time.Second}, nil - } - return r.failRunWithJob(ctx, &run, &job, resolved, "InvalidOutput", err.Error()) - } - return r.completeRun(ctx, &run, &job, resolved, results) - } +type triageJobMissingError struct { + name string + action wandbv2.TriageActionName +} - return r.markRunRunning(ctx, &run, &job, resolved) +func (e *triageJobMissingError) Error() string { + return fmt.Sprintf("Job %q disappeared before action %q completed", e.name, e.action) } func resolveTriageAction(application *wandbv2.Application, actionName string) (wandbv2.TriageActionSpec, error) { @@ -360,101 +433,166 @@ func resolvedTriageExecution( } } -func (r *TriageRunReconciler) markRunRunning( +func (r *TriageRunReconciler) actionStatus( ctx context.Context, - run *wandbv2.TriageRun, + action *resolvedTriageAction, job *batchv1.Job, - resolved *wandbv2.TriageResolvedExecution, -) (ctrl.Result, error) { - statusBefore := run.DeepCopy().Status - run.Status.Phase = wandbv2.TriageRunPhaseRunning - run.Status.ObservedGeneration = run.Generation - run.Status.JobRef = &corev1.LocalObjectReference{Name: job.Name} - run.Status.ResolvedExecution = resolved - if run.Status.StartedAt == nil { - switch { - case job.Status.StartTime != nil: - run.Status.StartedAt = job.Status.StartTime.DeepCopy() - case !job.CreationTimestamp.IsZero(): - run.Status.StartedAt = job.CreationTimestamp.DeepCopy() - default: - now := metav1.Now() - run.Status.StartedAt = &now + previous *wandbv2.TriageActionStatus, +) (wandbv2.TriageActionStatus, bool, string) { + status := wandbv2.TriageActionStatus{ + Action: action.name, + Phase: wandbv2.TriageRunPhaseRunning, + JobRef: &corev1.LocalObjectReference{Name: job.Name}, + ResolvedExecution: action.resolved, + StartedAt: triageActionStartTime(previous, job), + } + if jobFailed(job) { + status.Phase = wandbv2.TriageRunPhaseFailed + status.CompletedAt = triageCompletionTime(job) + message := jobConditionMessage(job, batchv1.JobFailed) + if message == "" { + message = fmt.Sprintf("Job %q failed", job.Name) + } + if results, err := r.collectTriageResults(ctx, job); err == nil { + status.Results = results + status.Summary = summarizeTriageResults(results) } + return status, false, fmt.Sprintf("action %q: %s", action.name, message) } - apimeta.SetStatusCondition(&run.Status.Conditions, metav1.Condition{ - Type: triageConditionSucceeded, - Status: metav1.ConditionUnknown, - ObservedGeneration: run.Generation, - Reason: "JobRunning", - Message: fmt.Sprintf("Job %q is running", job.Name), - }) - if apiequality.Semantic.DeepEqual(statusBefore, run.Status) { - return ctrl.Result{}, nil + if !jobComplete(job) { + return status, false, "" } - if err := r.Status().Update(ctx, run); err != nil { - return ctrl.Result{}, err + + results, err := r.collectTriageResults(ctx, job) + if err != nil { + var unavailable *triageOutputUnavailableError + if errors.As(err, &unavailable) { + return status, true, "" + } + status.Phase = wandbv2.TriageRunPhaseFailed + status.CompletedAt = triageCompletionTime(job) + return status, false, fmt.Sprintf("action %q: %s", action.name, err) } - return ctrl.Result{}, nil + status.Phase = wandbv2.TriageRunPhaseSucceeded + status.CompletedAt = triageCompletionTime(job) + status.Results = results + status.Summary = summarizeTriageResults(results) + return status, false, "" } -func (r *TriageRunReconciler) completeRun( - ctx context.Context, - run *wandbv2.TriageRun, - job *batchv1.Job, - resolved *wandbv2.TriageResolvedExecution, - results []wandbv2.TriageCheckResult, -) (ctrl.Result, error) { - run.Status.Phase = wandbv2.TriageRunPhaseSucceeded +func setAggregateTriageStatus(run *wandbv2.TriageRun, failureMessages []string) { run.Status.ObservedGeneration = run.Generation - run.Status.JobRef = &corev1.LocalObjectReference{Name: job.Name} - run.Status.ResolvedExecution = resolved - run.Status.Results = results - run.Status.Summary = summarizeTriageResults(results) - run.Status.StartedAt = triageStartTime(run, job) - run.Status.CompletedAt = triageCompletionTime(job) + run.Status.StartedAt = nil + run.Status.CompletedAt = nil + run.Status.Summary = nil + + completedActions := 0 + failedActions := 0 + var completedAt *metav1.Time + var summary wandbv2.TriageRunSummary + hasSummary := false + for i := range run.Status.ActionStatuses { + status := &run.Status.ActionStatuses[i] + if status.StartedAt != nil && + (run.Status.StartedAt == nil || status.StartedAt.Time.Before(run.Status.StartedAt.Time)) { + run.Status.StartedAt = status.StartedAt.DeepCopy() + } + if isTerminalTriagePhase(status.Phase) { + completedActions++ + if status.Phase == wandbv2.TriageRunPhaseFailed { + failedActions++ + } + if status.CompletedAt != nil && + (completedAt == nil || status.CompletedAt.Time.After(completedAt.Time)) { + completedAt = status.CompletedAt.DeepCopy() + } + } + if status.Summary != nil { + hasSummary = true + summary.Total += status.Summary.Total + summary.Pass += status.Summary.Pass + summary.Warn += status.Summary.Warn + summary.Fail += status.Summary.Fail + summary.Error += status.Summary.Error + } + } + if hasSummary { + setTriageSummarySeverity(&summary) + run.Status.Summary = &summary + } + + if completedActions < len(run.Status.ActionStatuses) { + run.Status.Phase = wandbv2.TriageRunPhaseRunning + apimeta.SetStatusCondition(&run.Status.Conditions, metav1.Condition{ + Type: triageConditionSucceeded, + Status: metav1.ConditionUnknown, + ObservedGeneration: run.Generation, + Reason: "ActionsRunning", + Message: fmt.Sprintf( + "%d of %d triage actions completed", completedActions, len(run.Status.ActionStatuses)), + }) + return + } + + run.Status.CompletedAt = completedAt + if run.Status.CompletedAt == nil { + now := metav1.Now() + run.Status.CompletedAt = &now + } + if failedActions > 0 { + run.Status.Phase = wandbv2.TriageRunPhaseFailed + message := strings.Join(failureMessages, "; ") + if message == "" { + message = fmt.Sprintf("%d triage actions failed", failedActions) + } + apimeta.SetStatusCondition(&run.Status.Conditions, metav1.Condition{ + Type: triageConditionSucceeded, + Status: metav1.ConditionFalse, + ObservedGeneration: run.Generation, + Reason: "ActionFailed", + Message: message, + }) + return + } + + run.Status.Phase = wandbv2.TriageRunPhaseSucceeded + totalResults := int32(0) + if run.Status.Summary != nil { + totalResults = run.Status.Summary.Total + } apimeta.SetStatusCondition(&run.Status.Conditions, metav1.Condition{ Type: triageConditionSucceeded, Status: metav1.ConditionTrue, ObservedGeneration: run.Generation, Reason: "ResultsCollected", - Message: fmt.Sprintf("Collected %d triage check results", len(results)), + Message: fmt.Sprintf( + "Collected %d triage check results from %d actions", totalResults, completedActions), }) - if err := r.Status().Update(ctx, run); err != nil { - return ctrl.Result{}, err - } - return ctrl.Result{}, nil } -func (r *TriageRunReconciler) failRun( - ctx context.Context, - run *wandbv2.TriageRun, - reason string, - message string, -) (ctrl.Result, error) { - return r.failRunWithJob(ctx, run, nil, run.Status.ResolvedExecution, reason, message) +func setTriageSummarySeverity(summary *wandbv2.TriageRunSummary) { + switch { + case summary.Error > 0: + summary.OverallSeverity = wandbv2.TriageSeverityError + case summary.Fail > 0: + summary.OverallSeverity = wandbv2.TriageSeverityFail + case summary.Warn > 0: + summary.OverallSeverity = wandbv2.TriageSeverityWarn + default: + summary.OverallSeverity = wandbv2.TriageSeverityPass + } } -func (r *TriageRunReconciler) failRunWithJob( +func (r *TriageRunReconciler) failRun( ctx context.Context, run *wandbv2.TriageRun, - job *batchv1.Job, - resolved *wandbv2.TriageResolvedExecution, reason string, message string, ) (ctrl.Result, error) { run.Status.Phase = wandbv2.TriageRunPhaseFailed run.Status.ObservedGeneration = run.Generation - run.Status.ResolvedExecution = resolved now := metav1.Now() run.Status.CompletedAt = &now - if job != nil { - run.Status.JobRef = &corev1.LocalObjectReference{Name: job.Name} - run.Status.StartedAt = triageStartTime(run, job) - if completedAt := triageCompletionTime(job); completedAt != nil { - run.Status.CompletedAt = completedAt - } - } apimeta.SetStatusCondition(&run.Status.Conditions, metav1.Condition{ Type: triageConditionSucceeded, Status: metav1.ConditionFalse, @@ -468,17 +606,21 @@ func (r *TriageRunReconciler) failRunWithJob( return ctrl.Result{}, nil } -func triageStartTime(run *wandbv2.TriageRun, job *batchv1.Job) *metav1.Time { +func triageActionStartTime( + previous *wandbv2.TriageActionStatus, + job *batchv1.Job, +) *metav1.Time { if job.Status.StartTime != nil { return job.Status.StartTime.DeepCopy() } - if run.Status.StartedAt != nil { - return run.Status.StartedAt.DeepCopy() + if previous != nil && previous.StartedAt != nil { + return previous.StartedAt.DeepCopy() } if !job.CreationTimestamp.IsZero() { return job.CreationTimestamp.DeepCopy() } - return nil + now := metav1.Now() + return &now } func triageCompletionTime(job *batchv1.Job) *metav1.Time { diff --git a/internal/controller/triagerun_controller_unit_test.go b/internal/controller/triagerun_controller_unit_test.go index 0cb094d5..3c79188b 100644 --- a/internal/controller/triagerun_controller_unit_test.go +++ b/internal/controller/triagerun_controller_unit_test.go @@ -31,10 +31,11 @@ func TestTriageRunCreatesBoundedJobFromApplication(t *testing.T) { WithStatusSubresource(&wandbv2.TriageRun{}, &batchv1.Job{}). WithObjects(run, application). Build() + logReader := &staticTriageLogReader{} reconciler := &TriageRunReconciler{ Client: fakeClient, Scheme: testScheme, - PodLogs: &staticTriageLogReader{}, + PodLogs: logReader, } _, err := reconciler.Reconcile(context.Background(), requestFor(run)) @@ -45,7 +46,7 @@ func TestTriageRunCreatesBoundedJobFromApplication(t *testing.T) { var job batchv1.Job if err := fakeClient.Get(context.Background(), types.NamespacedName{ Namespace: run.Namespace, - Name: "weave-check-triage", + Name: "weave-check-triage-0", }, &job); err != nil { t.Fatalf("get triage Job: %v", err) } @@ -107,12 +108,116 @@ func TestTriageRunCreatesBoundedJobFromApplication(t *testing.T) { if updatedRun.Status.Phase != wandbv2.TriageRunPhaseRunning { t.Fatalf("phase = %q, want Running", updatedRun.Status.Phase) } - if updatedRun.Status.JobRef == nil || updatedRun.Status.JobRef.Name != job.Name { - t.Fatalf("jobRef = %#v, want %q", updatedRun.Status.JobRef, job.Name) + if len(updatedRun.Status.ActionStatuses) != 1 { + t.Fatalf("action statuses = %d, want 1", len(updatedRun.Status.ActionStatuses)) } - if updatedRun.Status.ResolvedExecution == nil || - updatedRun.Status.ResolvedExecution.ContainerName != "weave-trace" { - t.Fatalf("resolved execution = %#v, want weave-trace container", updatedRun.Status.ResolvedExecution) + actionStatus := updatedRun.Status.ActionStatuses[0] + if actionStatus.Action != defaultTriageAction || + actionStatus.JobRef == nil || actionStatus.JobRef.Name != job.Name { + t.Fatalf("action status = %#v, want %q Job %q", actionStatus, defaultTriageAction, job.Name) + } + if actionStatus.ResolvedExecution == nil || + actionStatus.ResolvedExecution.ContainerName != "weave-trace" { + t.Fatalf("resolved execution = %#v, want weave-trace container", actionStatus.ResolvedExecution) + } +} + +func TestTriageRunCreatesOneJobPerSelectedAction(t *testing.T) { + t.Parallel() + + testScheme := newTriageTestScheme(t) + run := testTriageRun() + run.Spec.Actions = []wandbv2.TriageActionName{"default", "deep"} + application := testTriageApplication() + application.Spec.Triage.Actions["deep"] = wandbv2.TriageActionSpec{ + ContainerName: "weave-trace", + Args: []string{"python", "-m", "weave_triage", "deep"}, + } + fakeClient := fake.NewClientBuilder(). + WithScheme(testScheme). + WithStatusSubresource(&wandbv2.TriageRun{}, &batchv1.Job{}). + WithObjects(run, application). + Build() + logReader := &staticTriageLogReader{} + reconciler := &TriageRunReconciler{ + Client: fakeClient, + Scheme: testScheme, + PodLogs: logReader, + } + + if _, err := reconciler.Reconcile(context.Background(), requestFor(run)); err != nil { + t.Fatalf("reconcile multi-action TriageRun: %v", err) + } + + for index, action := range []string{"default", "deep"} { + var job batchv1.Job + name := fmt.Sprintf("weave-check-triage-%d", index) + if err := fakeClient.Get(context.Background(), types.NamespacedName{ + Namespace: run.Namespace, + Name: name, + }, &job); err != nil { + t.Fatalf("get Job for action %q: %v", action, err) + } + if job.Annotations[triageActionAnnotation] != action { + t.Fatalf("Job %q action annotation = %q, want %q", + job.Name, job.Annotations[triageActionAnnotation], action) + } + } + + var updatedRun wandbv2.TriageRun + if err := fakeClient.Get(context.Background(), client.ObjectKeyFromObject(run), &updatedRun); err != nil { + t.Fatalf("get updated TriageRun: %v", err) + } + if updatedRun.Status.Phase != wandbv2.TriageRunPhaseRunning || + len(updatedRun.Status.ActionStatuses) != 2 { + t.Fatalf("status = %#v, want two running actions", updatedRun.Status) + } + for i, action := range run.Spec.Actions { + if updatedRun.Status.ActionStatuses[i].Action != action { + t.Fatalf("action status %d = %q, want %q", + i, updatedRun.Status.ActionStatuses[i].Action, action) + } + } + + logReader.output = []byte(`{"name":"reachable","severity":"pass"}` + "\n") + for index := range run.Spec.Actions { + name := fmt.Sprintf("weave-check-triage-%d", index) + var job batchv1.Job + if err := fakeClient.Get(context.Background(), types.NamespacedName{ + Namespace: run.Namespace, + Name: name, + }, &job); err != nil { + t.Fatalf("get Job %q for completion: %v", name, err) + } + job.Status.Conditions = []batchv1.JobCondition{{ + Type: batchv1.JobComplete, Status: corev1.ConditionTrue, + }} + if err := fakeClient.Status().Update(context.Background(), &job); err != nil { + t.Fatalf("mark Job %q complete: %v", name, err) + } + pod := &corev1.Pod{ObjectMeta: metav1.ObjectMeta{ + Name: name + "-pod", + Namespace: run.Namespace, + Labels: map[string]string{"batch.kubernetes.io/job-name": name}, + }} + if err := fakeClient.Create(context.Background(), pod); err != nil { + t.Fatalf("create pod for Job %q: %v", name, err) + } + } + if _, err := reconciler.Reconcile(context.Background(), requestFor(run)); err != nil { + t.Fatalf("reconcile completed multi-action TriageRun: %v", err) + } + if err := fakeClient.Get(context.Background(), client.ObjectKeyFromObject(run), &updatedRun); err != nil { + t.Fatalf("get completed TriageRun: %v", err) + } + if updatedRun.Status.Phase != wandbv2.TriageRunPhaseSucceeded || + updatedRun.Status.Summary == nil || updatedRun.Status.Summary.Total != 2 { + t.Fatalf("completed status = %#v, want two aggregated results", updatedRun.Status) + } + for _, status := range updatedRun.Status.ActionStatuses { + if status.Phase != wandbv2.TriageRunPhaseSucceeded || len(status.Results) != 1 { + t.Fatalf("action status = %#v, want one successful result", status) + } } } @@ -122,12 +227,15 @@ func TestTriageRunCollectsFailedCheckAsSuccessfulExecution(t *testing.T) { testScheme := newTriageTestScheme(t) run := testTriageRun() run.Status.Phase = wandbv2.TriageRunPhaseRunning - run.Status.JobRef = &corev1.LocalObjectReference{Name: "weave-check-triage"} + run.Status.ActionStatuses = []wandbv2.TriageActionStatus{{ + Action: defaultTriageAction, + JobRef: &corev1.LocalObjectReference{Name: "weave-check-triage-0"}, + }} application := testTriageApplication() action := application.Spec.Triage.Actions[defaultTriageAction] source := &application.Spec.PodTemplate.Spec.Containers[0] job := buildTriageJob( - run, application, source, action, defaultTriageTimeoutSeconds, "weave-check-triage") + run, application, source, action, defaultTriageTimeoutSeconds, "weave-check-triage-0") if err := controllerutil.SetControllerReference(run, job, testScheme); err != nil { t.Fatalf("set Job owner: %v", err) } @@ -174,8 +282,9 @@ func TestTriageRunCollectsFailedCheckAsSuccessfulExecution(t *testing.T) { updatedRun.Status.Summary.Fail != 1 { t.Fatalf("summary = %#v, want one failed diagnostic check", updatedRun.Status.Summary) } - if len(updatedRun.Status.Results) != 2 { - t.Fatalf("results = %d, want 2", len(updatedRun.Status.Results)) + if len(updatedRun.Status.ActionStatuses) != 1 || + len(updatedRun.Status.ActionStatuses[0].Results) != 2 { + t.Fatalf("action statuses = %#v, want 2 results", updatedRun.Status.ActionStatuses) } } @@ -261,7 +370,7 @@ func testTriageRun() *wandbv2.TriageRun { }, Spec: wandbv2.TriageRunSpec{ ApplicationRef: wandbv2.TriageApplicationReference{Name: "weave-trace"}, - Action: defaultTriageAction, + Actions: []wandbv2.TriageActionName{defaultTriageAction}, }, } } From a205da8b91d6aa8f1d1493678aa5b67e31b968b8 Mon Sep 17 00:00:00 2001 From: Aravind Warrier Date: Tue, 4 Aug 2026 12:29:14 -0400 Subject: [PATCH 3/4] refactor(controller): resolve structured triage actions --- api/v2/application_types.go | 46 ++- api/v2/zz_generated.deepcopy.go | 45 +-- .../bases/apps.wandb.com_applications.yaml | 276 +++++++++--------- config/samples/apps_v2_application.yaml | 9 +- internal/controller/reconciler/triage.go | 40 +-- internal/controller/reconciler/triage_test.go | 53 ++-- internal/controller/triagerun_controller.go | 69 +++-- .../triagerun_controller_unit_test.go | 79 +++-- .../operator/apps.wandb.com_applications.yaml | 276 +++++++++--------- pkg/wandb/manifest/manifest.go | 17 +- 10 files changed, 507 insertions(+), 403 deletions(-) diff --git a/api/v2/application_types.go b/api/v2/application_types.go index e8a28352..731da7e4 100644 --- a/api/v2/application_types.go +++ b/api/v2/application_types.go @@ -77,20 +77,10 @@ type ApplicationSpec struct { HTTPRouteTemplate *HTTPRouteTemplateSpec `json:"httpRouteTemplate,omitempty"` } -// ApplicationTriageSpec contains the diagnostic actions exposed by an -// Application. The Application remains the source of runtime configuration; -// TriageRun only selects one of these actions. +// ApplicationTriageSpec contains the shared diagnostic runner and the actions +// exposed by an Application. TriageRun selects actions by name; the controller +// appends each selected name to the resolved runner arguments. type ApplicationTriageSpec struct { - // Actions maps stable action names to their execution overrides. - // +kubebuilder:validation:MinProperties=1 - // +kubebuilder:validation:MaxProperties=16 - Actions map[string]TriageActionSpec `json:"actions"` -} - -// TriageActionSpec is a compact override applied to one container from the -// Application pod template. Image, identity, environment, volumes, and -// scheduling settings are inherited from the Application. -type TriageActionSpec struct { // ContainerName selects a container from the Application pod template. It // may be omitted when the Application has exactly one container. // +optional @@ -102,7 +92,8 @@ type TriageActionSpec struct { // +optional Command []string `json:"command,omitempty"` - // Args replaces the selected container's arguments when non-empty. + // Args replaces the selected container's arguments when non-empty. The + // controller appends the selected action name and then that action's Args. // +kubebuilder:validation:MinItems=1 // +kubebuilder:validation:MaxItems=64 // +optional @@ -126,6 +117,33 @@ type TriageActionSpec struct { // +kubebuilder:validation:Maximum=3600 // +optional TimeoutSeconds int64 `json:"timeoutSeconds,omitempty"` + + // Actions lists the stable action names and metadata exposed to callers. + // +kubebuilder:validation:MinItems=1 + // +kubebuilder:validation:MaxItems=16 + // +listType=map + // +listMapKey=name + Actions []TriageActionSpec `json:"actions"` +} + +// TriageActionSpec describes one action exposed by the shared diagnostic +// runner. Execution identity and resource settings remain on the parent +// ApplicationTriageSpec so every action uses the same bounded runtime. +type TriageActionSpec struct { + // Name is the stable identifier selected by TriageRun and passed to the + // diagnostic runner as its next argument. + Name TriageActionName `json:"name"` + + // Description is human-readable help shown by clients such as Watchtower. + // +kubebuilder:validation:MaxLength=512 + // +optional + Description string `json:"description,omitempty"` + + // Args are appended after the action name when starting the diagnostic + // runner. They are suitable for action-specific flags, not executable paths. + // +kubebuilder:validation:MaxItems=32 + // +optional + Args []string `json:"args,omitempty"` } // HTTPRouteTemplateSpec contains the fields needed to build a Gateway API HTTPRoute. diff --git a/api/v2/zz_generated.deepcopy.go b/api/v2/zz_generated.deepcopy.go index 514b095a..53a6af4e 100644 --- a/api/v2/zz_generated.deepcopy.go +++ b/api/v2/zz_generated.deepcopy.go @@ -239,11 +239,33 @@ func (in *ApplicationStatus) DeepCopy() *ApplicationStatus { // DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. func (in *ApplicationTriageSpec) DeepCopyInto(out *ApplicationTriageSpec) { *out = *in + if in.Command != nil { + in, out := &in.Command, &out.Command + *out = make([]string, len(*in)) + copy(*out, *in) + } + if in.Args != nil { + in, out := &in.Args, &out.Args + *out = make([]string, len(*in)) + copy(*out, *in) + } + if in.Env != nil { + in, out := &in.Env, &out.Env + *out = make([]v1.EnvVar, len(*in)) + for i := range *in { + (*in)[i].DeepCopyInto(&(*out)[i]) + } + } + if in.Resources != nil { + in, out := &in.Resources, &out.Resources + *out = new(v1.ResourceRequirements) + (*in).DeepCopyInto(*out) + } if in.Actions != nil { in, out := &in.Actions, &out.Actions - *out = make(map[string]TriageActionSpec, len(*in)) - for key, val := range *in { - (*out)[key] = *val.DeepCopy() + *out = make([]TriageActionSpec, len(*in)) + for i := range *in { + (*in)[i].DeepCopyInto(&(*out)[i]) } } } @@ -1524,28 +1546,11 @@ func (in *TriageActionReference) DeepCopy() *TriageActionReference { // DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. func (in *TriageActionSpec) DeepCopyInto(out *TriageActionSpec) { *out = *in - if in.Command != nil { - in, out := &in.Command, &out.Command - *out = make([]string, len(*in)) - copy(*out, *in) - } if in.Args != nil { in, out := &in.Args, &out.Args *out = make([]string, len(*in)) copy(*out, *in) } - if in.Env != nil { - in, out := &in.Env, &out.Env - *out = make([]v1.EnvVar, len(*in)) - for i := range *in { - (*in)[i].DeepCopyInto(&(*out)[i]) - } - } - if in.Resources != nil { - in, out := &in.Resources, &out.Resources - *out = new(v1.ResourceRequirements) - (*in).DeepCopyInto(*out) - } } // DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new TriageActionSpec. diff --git a/config/crd/bases/apps.wandb.com_applications.yaml b/config/crd/bases/apps.wandb.com_applications.yaml index 5cd5c9dd..3450d383 100644 --- a/config/crd/bases/apps.wandb.com_applications.yaml +++ b/config/crd/bases/apps.wandb.com_applications.yaml @@ -12804,148 +12804,164 @@ spec: triage: properties: actions: - additionalProperties: + items: properties: args: items: type: string - maxItems: 64 - minItems: 1 + maxItems: 32 type: array - command: - items: - type: string - maxItems: 64 - minItems: 1 - type: array - containerName: + description: + maxLength: 512 type: string - env: - items: - properties: - name: - type: string - value: - type: string - valueFrom: - properties: - configMapKeyRef: - properties: - key: - type: string - name: - default: "" - type: string - optional: - type: boolean - required: - - key - type: object - x-kubernetes-map-type: atomic - fieldRef: - properties: - apiVersion: - type: string - fieldPath: - type: string - required: - - fieldPath - type: object - x-kubernetes-map-type: atomic - fileKeyRef: - properties: - key: - type: string - optional: - default: false - type: boolean - path: - type: string - volumeName: - type: string - required: - - key - - path - - volumeName - type: object - x-kubernetes-map-type: atomic - resourceFieldRef: - properties: - containerName: - type: string - divisor: - anyOf: - - type: integer - - type: string - pattern: ^(\+|-)?(([0-9]+(\.[0-9]*)?)|(\.[0-9]+))(([KMGTPE]i)|[numkMGTPE]|([eE](\+|-)?(([0-9]+(\.[0-9]*)?)|(\.[0-9]+))))?$ - x-kubernetes-int-or-string: true - resource: - type: string - required: - - resource - type: object - x-kubernetes-map-type: atomic - secretKeyRef: - properties: - key: - type: string - name: - default: "" - type: string - optional: - type: boolean - required: - - key - type: object - x-kubernetes-map-type: atomic - type: object - required: - - name - type: object - maxItems: 128 - type: array - resources: + name: + minLength: 1 + type: string + required: + - name + type: object + maxItems: 16 + minItems: 1 + type: array + x-kubernetes-list-map-keys: + - name + x-kubernetes-list-type: map + args: + items: + type: string + maxItems: 64 + minItems: 1 + type: array + command: + items: + type: string + maxItems: 64 + minItems: 1 + type: array + containerName: + type: string + env: + items: + properties: + name: + type: string + value: + type: string + valueFrom: properties: - claims: - items: - properties: - name: - type: string - request: - type: string - required: - - name - type: object - type: array - x-kubernetes-list-map-keys: - - name - x-kubernetes-list-type: map - limits: - additionalProperties: - anyOf: - - type: integer - - type: string - pattern: ^(\+|-)?(([0-9]+(\.[0-9]*)?)|(\.[0-9]+))(([KMGTPE]i)|[numkMGTPE]|([eE](\+|-)?(([0-9]+(\.[0-9]*)?)|(\.[0-9]+))))?$ - x-kubernetes-int-or-string: true + configMapKeyRef: + properties: + key: + type: string + name: + default: "" + type: string + optional: + type: boolean + required: + - key type: object - requests: - additionalProperties: - anyOf: - - type: integer - - type: string - pattern: ^(\+|-)?(([0-9]+(\.[0-9]*)?)|(\.[0-9]+))(([KMGTPE]i)|[numkMGTPE]|([eE](\+|-)?(([0-9]+(\.[0-9]*)?)|(\.[0-9]+))))?$ - x-kubernetes-int-or-string: true + x-kubernetes-map-type: atomic + fieldRef: + properties: + apiVersion: + type: string + fieldPath: + type: string + required: + - fieldPath + type: object + x-kubernetes-map-type: atomic + fileKeyRef: + properties: + key: + type: string + optional: + default: false + type: boolean + path: + type: string + volumeName: + type: string + required: + - key + - path + - volumeName + type: object + x-kubernetes-map-type: atomic + resourceFieldRef: + properties: + containerName: + type: string + divisor: + anyOf: + - type: integer + - type: string + pattern: ^(\+|-)?(([0-9]+(\.[0-9]*)?)|(\.[0-9]+))(([KMGTPE]i)|[numkMGTPE]|([eE](\+|-)?(([0-9]+(\.[0-9]*)?)|(\.[0-9]+))))?$ + x-kubernetes-int-or-string: true + resource: + type: string + required: + - resource + type: object + x-kubernetes-map-type: atomic + secretKeyRef: + properties: + key: + type: string + name: + default: "" + type: string + optional: + type: boolean + required: + - key type: object + x-kubernetes-map-type: atomic type: object - timeoutSeconds: - format: int64 - maximum: 3600 - minimum: 1 - type: integer + required: + - name type: object - maxProperties: 16 - minProperties: 1 + maxItems: 128 + type: array + resources: + properties: + claims: + items: + properties: + name: + type: string + request: + type: string + required: + - name + type: object + type: array + x-kubernetes-list-map-keys: + - name + x-kubernetes-list-type: map + limits: + additionalProperties: + anyOf: + - type: integer + - type: string + pattern: ^(\+|-)?(([0-9]+(\.[0-9]*)?)|(\.[0-9]+))(([KMGTPE]i)|[numkMGTPE]|([eE](\+|-)?(([0-9]+(\.[0-9]*)?)|(\.[0-9]+))))?$ + x-kubernetes-int-or-string: true + type: object + requests: + additionalProperties: + anyOf: + - type: integer + - type: string + pattern: ^(\+|-)?(([0-9]+(\.[0-9]*)?)|(\.[0-9]+))(([KMGTPE]i)|[numkMGTPE]|([eE](\+|-)?(([0-9]+(\.[0-9]*)?)|(\.[0-9]+))))?$ + x-kubernetes-int-or-string: true + type: object type: object + timeoutSeconds: + format: int64 + maximum: 3600 + minimum: 1 + type: integer required: - actions type: object diff --git a/config/samples/apps_v2_application.yaml b/config/samples/apps_v2_application.yaml index 181ebc5f..af7f7a61 100644 --- a/config/samples/apps_v2_application.yaml +++ b/config/samples/apps_v2_application.yaml @@ -17,11 +17,12 @@ spec: image: nginx:1.14.2 triage: + containerName: my-app + command: ["/bin/echo"] + timeoutSeconds: 60 actions: - default: - containerName: my-app - command: ["nginx", "-t"] - timeoutSeconds: 60 + - name: default + description: Run the default application diagnostic resources: requests: cpu: 25m diff --git a/internal/controller/reconciler/triage.go b/internal/controller/reconciler/triage.go index 1bb16721..98145e1b 100644 --- a/internal/controller/reconciler/triage.go +++ b/internal/controller/reconciler/triage.go @@ -11,24 +11,28 @@ func resolveApplicationTriage(triage *serverManifest.ApplicationTriage) *apiv2.A return nil } - actions := make(map[string]apiv2.TriageActionSpec, len(triage.Actions)) - for name, action := range triage.Actions { - env := make([]corev1.EnvVar, len(action.Env)) - for i := range action.Env { - env[i] = *action.Env[i].DeepCopy() - } - resolved := apiv2.TriageActionSpec{ - ContainerName: action.ContainerName, - Command: append([]string(nil), action.Command...), - Args: append([]string(nil), action.Args...), - Env: env, - TimeoutSeconds: action.TimeoutSeconds, - } - if action.Resources != nil { - resolved.Resources = action.Resources.DeepCopy() + env := make([]corev1.EnvVar, len(triage.Env)) + for i := range triage.Env { + env[i] = *triage.Env[i].DeepCopy() + } + actions := make([]apiv2.TriageActionSpec, len(triage.Actions)) + for i := range triage.Actions { + actions[i] = apiv2.TriageActionSpec{ + Name: apiv2.TriageActionName(triage.Actions[i].Name), + Description: triage.Actions[i].Description, + Args: append([]string(nil), triage.Actions[i].Args...), } - actions[name] = resolved } - - return &apiv2.ApplicationTriageSpec{Actions: actions} + resolved := &apiv2.ApplicationTriageSpec{ + ContainerName: triage.ContainerName, + Command: append([]string(nil), triage.Command...), + Args: append([]string(nil), triage.Args...), + Env: env, + TimeoutSeconds: triage.TimeoutSeconds, + Actions: actions, + } + if triage.Resources != nil { + resolved.Resources = triage.Resources.DeepCopy() + } + return resolved } diff --git a/internal/controller/reconciler/triage_test.go b/internal/controller/reconciler/triage_test.go index 93f8c682..423f46bd 100644 --- a/internal/controller/reconciler/triage_test.go +++ b/internal/controller/reconciler/triage_test.go @@ -16,38 +16,45 @@ func TestResolveApplicationTriageCopiesCompactAction(t *testing.T) { }, } input := &serverManifest.ApplicationTriage{ - Actions: map[string]serverManifest.TriageAction{ - "default": { - ContainerName: "weave-trace", - Args: []string{"python", "-m", "weave_triage", "run-all", "--stream"}, - Env: []corev1.EnvVar{{ - Name: "PYTHONPATH", - Value: "/weave/src", - }}, - Resources: resources, - TimeoutSeconds: 600, - }, - }, + ContainerName: "weave-trace", + Args: []string{"python", "-m", "weave_triage"}, + Env: []corev1.EnvVar{{ + Name: "PYTHONPATH", + Value: "/weave/src", + }}, + Resources: resources, + TimeoutSeconds: 600, + Actions: []serverManifest.TriageAction{{ + Name: "default", + Description: "Run all diagnostics", + Args: []string{"--verbose"}, + }}, } resolved := resolveApplicationTriage(input) - action := resolved.Actions["default"] - if action.ContainerName != "weave-trace" { - t.Fatalf("containerName = %q", action.ContainerName) + action := resolved.Actions[0] + if resolved.ContainerName != "weave-trace" { + t.Fatalf("containerName = %q", resolved.ContainerName) + } + if len(resolved.Args) != 3 || resolved.Args[0] != "python" { + t.Fatalf("runner args = %#v", resolved.Args) } - if len(action.Args) != 5 || action.Args[4] != "--stream" { - t.Fatalf("args = %#v", action.Args) + if resolved.Resources == nil || resolved.Resources.Requests.Memory().String() != "128Mi" { + t.Fatalf("resources = %#v", resolved.Resources) } - if action.Resources == nil || action.Resources.Requests.Memory().String() != "128Mi" { - t.Fatalf("resources = %#v", action.Resources) + if resolved.TimeoutSeconds != 600 { + t.Fatalf("timeoutSeconds = %d", resolved.TimeoutSeconds) } - if action.TimeoutSeconds != 600 { - t.Fatalf("timeoutSeconds = %d", action.TimeoutSeconds) + if action.Name != "default" || action.Description != "Run all diagnostics" || + len(action.Args) != 1 || action.Args[0] != "--verbose" { + t.Fatalf("action = %#v", action) } - input.Actions["default"].Args[0] = "mutated" + input.Args[0] = "mutated" + input.Actions[0].Args[0] = "mutated" resources.Requests[corev1.ResourceMemory] = resource.MustParse("1Gi") - if action.Args[0] != "python" || action.Resources.Requests.Memory().String() != "128Mi" { + if resolved.Args[0] != "python" || action.Args[0] != "--verbose" || + resolved.Resources.Requests.Memory().String() != "128Mi" { t.Fatal("resolved action aliases mutable manifest data") } } diff --git a/internal/controller/triagerun_controller.go b/internal/controller/triagerun_controller.go index 7d1fed5a..99c97aa9 100644 --- a/internal/controller/triagerun_controller.go +++ b/internal/controller/triagerun_controller.go @@ -185,6 +185,7 @@ func (r *TriageRunReconciler) Reconcile(ctx context.Context, req ctrl.Request) ( type resolvedTriageAction struct { name wandbv2.TriageActionName + runner *wandbv2.ApplicationTriageSpec spec wandbv2.TriageActionSpec source *corev1.Container timeoutSeconds int64 @@ -200,25 +201,29 @@ func resolveRequestedTriageActions( return nil, errors.New("at least one triage action is required") } actions := make([]resolvedTriageAction, 0, len(run.Spec.Actions)) - for i, actionName := range run.Spec.Actions { - spec, err := resolveTriageAction(application, string(actionName)) + for i, actionRef := range run.Spec.Actions { + actionName := actionRef.Name + spec, err := resolveTriageAction(application, actionName) if err != nil { return nil, err } - source, err := selectTriageContainer(application, spec.ContainerName) + runner := application.Spec.Triage + source, err := selectTriageContainer(application, runner.ContainerName) if err != nil { return nil, err } - timeoutSeconds := spec.TimeoutSeconds + timeoutSeconds := runner.TimeoutSeconds if timeoutSeconds == 0 { timeoutSeconds = defaultTriageTimeoutSeconds } actions = append(actions, resolvedTriageAction{ name: actionName, + runner: runner, spec: spec, source: source, timeoutSeconds: timeoutSeconds, - resolved: resolvedTriageExecution(application, source, spec, timeoutSeconds), + resolved: resolvedTriageExecution( + application, source, runner, spec, timeoutSeconds), jobName: common.FitDefaultInfraName( run.Name, "-triage-"+strconv.Itoa(i), 63), }) @@ -256,8 +261,8 @@ func (r *TriageRunReconciler) getOrCreateTriageJob( name: previous.JobRef.Name, action: action.name, } } - job = *buildTriageJob( - run, application, action.source, action.spec, action.timeoutSeconds, action.jobName) + job = *buildTriageJob(run, application, action.source, action.runner, + action.spec, action.timeoutSeconds, action.jobName) job.Annotations = map[string]string{triageActionAnnotation: string(action.name)} if err := controllerutil.SetControllerReference(run, &job, r.Scheme); err != nil { return nil, err @@ -286,20 +291,25 @@ func (e *triageJobMissingError) Error() string { return fmt.Sprintf("Job %q disappeared before action %q completed", e.name, e.action) } -func resolveTriageAction(application *wandbv2.Application, actionName string) (wandbv2.TriageActionSpec, error) { +func resolveTriageAction( + application *wandbv2.Application, + actionName wandbv2.TriageActionName, +) (wandbv2.TriageActionSpec, error) { if application.Spec.Triage == nil { return wandbv2.TriageActionSpec{}, fmt.Errorf("Application %q does not declare triage actions", application.Name) } - action, ok := application.Spec.Triage.Actions[actionName] - if !ok { + if len(application.Spec.Triage.Command) == 0 && len(application.Spec.Triage.Args) == 0 { return wandbv2.TriageActionSpec{}, fmt.Errorf( - "Application %q does not declare triage action %q", application.Name, actionName) + "Application %q triage runner must override command or args", application.Name) } - if len(action.Command) == 0 && len(action.Args) == 0 { - return wandbv2.TriageActionSpec{}, fmt.Errorf( - "triage action %q must override command or args", actionName) + for i := range application.Spec.Triage.Actions { + action := application.Spec.Triage.Actions[i] + if action.Name == actionName { + return action, nil + } } - return action, nil + return wandbv2.TriageActionSpec{}, fmt.Errorf( + "Application %q does not declare triage action %q", application.Name, actionName) } func selectTriageContainer(application *wandbv2.Application, name string) (*corev1.Container, error) { @@ -324,6 +334,7 @@ func buildTriageJob( run *wandbv2.TriageRun, application *wandbv2.Application, source *corev1.Container, + runner *wandbv2.ApplicationTriageSpec, action wandbv2.TriageActionSpec, timeoutSeconds int64, jobName string, @@ -334,7 +345,7 @@ func buildTriageJob( podSpec.InitContainers = nil podSpec.EphemeralContainers = nil podSpec.ReadinessGates = nil - podSpec.Containers = []corev1.Container{buildTriageContainer(source, action)} + podSpec.Containers = []corev1.Container{buildTriageContainer(source, runner, action)} labels := map[string]string{ triageRunLabel: common.FitDefaultInfraName(run.Name, "", 63), @@ -357,7 +368,11 @@ func buildTriageJob( } } -func buildTriageContainer(source *corev1.Container, action wandbv2.TriageActionSpec) corev1.Container { +func buildTriageContainer( + source *corev1.Container, + runner *wandbv2.ApplicationTriageSpec, + action wandbv2.TriageActionSpec, +) corev1.Container { container := corev1.Container{ Name: triageContainerName, Image: source.Image, @@ -366,7 +381,7 @@ func buildTriageContainer(source *corev1.Container, action wandbv2.TriageActionS Args: append([]string(nil), source.Args...), WorkingDir: source.WorkingDir, EnvFrom: append([]corev1.EnvFromSource(nil), source.EnvFrom...), - Env: mergeTriageEnv(source.Env, action.Env), + Env: mergeTriageEnv(source.Env, runner.Env), Resources: defaultTriageResources(), VolumeMounts: append([]corev1.VolumeMount(nil), source.VolumeMounts...), VolumeDevices: append([]corev1.VolumeDevice(nil), source.VolumeDevices...), @@ -374,14 +389,17 @@ func buildTriageContainer(source *corev1.Container, action wandbv2.TriageActionS TerminationMessagePath: source.TerminationMessagePath, TerminationMessagePolicy: source.TerminationMessagePolicy, } - if len(action.Command) > 0 { - container.Command = append([]string(nil), action.Command...) + if len(runner.Command) > 0 { + container.Command = append([]string(nil), runner.Command...) + container.Args = nil } - if len(action.Args) > 0 { - container.Args = append([]string(nil), action.Args...) + if len(runner.Args) > 0 { + container.Args = append([]string(nil), runner.Args...) } - if action.Resources != nil { - container.Resources = *action.Resources.DeepCopy() + container.Args = append(container.Args, string(action.Name)) + container.Args = append(container.Args, action.Args...) + if runner.Resources != nil { + container.Resources = *runner.Resources.DeepCopy() } return container } @@ -419,10 +437,11 @@ func defaultTriageResources() corev1.ResourceRequirements { func resolvedTriageExecution( application *wandbv2.Application, source *corev1.Container, + runner *wandbv2.ApplicationTriageSpec, action wandbv2.TriageActionSpec, timeoutSeconds int64, ) *wandbv2.TriageResolvedExecution { - container := buildTriageContainer(source, action) + container := buildTriageContainer(source, runner, action) return &wandbv2.TriageResolvedExecution{ ApplicationGeneration: application.Generation, ContainerName: source.Name, diff --git a/internal/controller/triagerun_controller_unit_test.go b/internal/controller/triagerun_controller_unit_test.go index 3c79188b..07507e52 100644 --- a/internal/controller/triagerun_controller_unit_test.go +++ b/internal/controller/triagerun_controller_unit_test.go @@ -82,7 +82,8 @@ func TestTriageRunCreatesBoundedJobFromApplication(t *testing.T) { ) >= 0 { t.Fatal("triage memory request must be smaller than the parent application request") } - if len(container.Args) != 4 || container.Args[0] != "python" { + if len(container.Args) != 4 || container.Args[0] != "python" || + container.Args[3] != defaultTriageAction { t.Fatalf("args = %#v, want triage action args", container.Args) } if got := envValue(container.Env, "PYTHONPATH"); got != "/weave/src" { @@ -127,12 +128,16 @@ func TestTriageRunCreatesOneJobPerSelectedAction(t *testing.T) { testScheme := newTriageTestScheme(t) run := testTriageRun() - run.Spec.Actions = []wandbv2.TriageActionName{"default", "deep"} - application := testTriageApplication() - application.Spec.Triage.Actions["deep"] = wandbv2.TriageActionSpec{ - ContainerName: "weave-trace", - Args: []string{"python", "-m", "weave_triage", "deep"}, + run.Spec.Actions = []wandbv2.TriageActionReference{ + {Name: "default"}, + {Name: "deep"}, } + application := testTriageApplication() + application.Spec.Triage.Actions = append(application.Spec.Triage.Actions, + wandbv2.TriageActionSpec{ + Name: "deep", + Description: "Run deeper diagnostics", + }) fakeClient := fake.NewClientBuilder(). WithScheme(testScheme). WithStatusSubresource(&wandbv2.TriageRun{}, &batchv1.Job{}). @@ -162,6 +167,10 @@ func TestTriageRunCreatesOneJobPerSelectedAction(t *testing.T) { t.Fatalf("Job %q action annotation = %q, want %q", job.Name, job.Annotations[triageActionAnnotation], action) } + args := job.Spec.Template.Spec.Containers[0].Args + if len(args) != 4 || args[3] != action { + t.Fatalf("Job %q args = %#v, want selected action %q appended", job.Name, args, action) + } } var updatedRun wandbv2.TriageRun @@ -173,9 +182,9 @@ func TestTriageRunCreatesOneJobPerSelectedAction(t *testing.T) { t.Fatalf("status = %#v, want two running actions", updatedRun.Status) } for i, action := range run.Spec.Actions { - if updatedRun.Status.ActionStatuses[i].Action != action { + if updatedRun.Status.ActionStatuses[i].Action != action.Name { t.Fatalf("action status %d = %q, want %q", - i, updatedRun.Status.ActionStatuses[i].Action, action) + i, updatedRun.Status.ActionStatuses[i].Action, action.Name) } } @@ -232,10 +241,11 @@ func TestTriageRunCollectsFailedCheckAsSuccessfulExecution(t *testing.T) { JobRef: &corev1.LocalObjectReference{Name: "weave-check-triage-0"}, }} application := testTriageApplication() - action := application.Spec.Triage.Actions[defaultTriageAction] + action := application.Spec.Triage.Actions[0] source := &application.Spec.PodTemplate.Spec.Containers[0] job := buildTriageJob( - run, application, source, action, defaultTriageTimeoutSeconds, "weave-check-triage-0") + run, application, source, application.Spec.Triage, action, + defaultTriageTimeoutSeconds, "weave-check-triage-0") if err := controllerutil.SetControllerReference(run, job, testScheme); err != nil { t.Fatalf("set Job owner: %v", err) } @@ -308,22 +318,25 @@ func TestManifestTriageActionDecodes(t *testing.T) { applications: weave-trace: triage: + containerName: weave-trace + args: [python, -m, weave_triage] + timeoutSeconds: 600 + resources: + requests: + cpu: 100m + memory: 128Mi actions: - default: - containerName: weave-trace - args: [python, -m, weave_triage, run-all, --stream] - timeoutSeconds: 600 - resources: - requests: - cpu: 100m - memory: 128Mi + + - name: default + description: Run all diagnostics `) if err := yaml.Unmarshal(input, &decoded); err != nil { t.Fatalf("decode manifest triage action: %v", err) } - action := decoded.Applications["weave-trace"].Triage.Actions["default"] - if action.ContainerName != "weave-trace" || action.TimeoutSeconds != 600 { - t.Fatalf("decoded action = %#v", action) + triage := decoded.Applications["weave-trace"].Triage + if triage.ContainerName != "weave-trace" || triage.TimeoutSeconds != 600 || + len(triage.Actions) != 1 || triage.Actions[0].Name != defaultTriageAction { + t.Fatalf("decoded triage = %#v", triage) } } @@ -370,7 +383,9 @@ func testTriageRun() *wandbv2.TriageRun { }, Spec: wandbv2.TriageRunSpec{ ApplicationRef: wandbv2.TriageApplicationReference{Name: "weave-trace"}, - Actions: []wandbv2.TriageActionName{defaultTriageAction}, + Actions: []wandbv2.TriageActionReference{{ + Name: defaultTriageAction, + }}, }, } } @@ -420,16 +435,16 @@ func testTriageApplication() *wandbv2.Application { }, }, Triage: &wandbv2.ApplicationTriageSpec{ - Actions: map[string]wandbv2.TriageActionSpec{ - defaultTriageAction: { - ContainerName: "weave-trace", - Args: []string{"python", "-m", "weave_triage", "run-all"}, - Env: []corev1.EnvVar{{ - Name: "PYTHONPATH", - Value: "/weave/src", - }}, - }, - }, + ContainerName: "weave-trace", + Args: []string{"python", "-m", "weave_triage"}, + Env: []corev1.EnvVar{{ + Name: "PYTHONPATH", + Value: "/weave/src", + }}, + Actions: []wandbv2.TriageActionSpec{{ + Name: defaultTriageAction, + Description: "Run all diagnostics", + }}, }, }, } diff --git a/internal/crdinstaller/crds/operator/apps.wandb.com_applications.yaml b/internal/crdinstaller/crds/operator/apps.wandb.com_applications.yaml index 5cd5c9dd..3450d383 100644 --- a/internal/crdinstaller/crds/operator/apps.wandb.com_applications.yaml +++ b/internal/crdinstaller/crds/operator/apps.wandb.com_applications.yaml @@ -12804,148 +12804,164 @@ spec: triage: properties: actions: - additionalProperties: + items: properties: args: items: type: string - maxItems: 64 - minItems: 1 + maxItems: 32 type: array - command: - items: - type: string - maxItems: 64 - minItems: 1 - type: array - containerName: + description: + maxLength: 512 type: string - env: - items: - properties: - name: - type: string - value: - type: string - valueFrom: - properties: - configMapKeyRef: - properties: - key: - type: string - name: - default: "" - type: string - optional: - type: boolean - required: - - key - type: object - x-kubernetes-map-type: atomic - fieldRef: - properties: - apiVersion: - type: string - fieldPath: - type: string - required: - - fieldPath - type: object - x-kubernetes-map-type: atomic - fileKeyRef: - properties: - key: - type: string - optional: - default: false - type: boolean - path: - type: string - volumeName: - type: string - required: - - key - - path - - volumeName - type: object - x-kubernetes-map-type: atomic - resourceFieldRef: - properties: - containerName: - type: string - divisor: - anyOf: - - type: integer - - type: string - pattern: ^(\+|-)?(([0-9]+(\.[0-9]*)?)|(\.[0-9]+))(([KMGTPE]i)|[numkMGTPE]|([eE](\+|-)?(([0-9]+(\.[0-9]*)?)|(\.[0-9]+))))?$ - x-kubernetes-int-or-string: true - resource: - type: string - required: - - resource - type: object - x-kubernetes-map-type: atomic - secretKeyRef: - properties: - key: - type: string - name: - default: "" - type: string - optional: - type: boolean - required: - - key - type: object - x-kubernetes-map-type: atomic - type: object - required: - - name - type: object - maxItems: 128 - type: array - resources: + name: + minLength: 1 + type: string + required: + - name + type: object + maxItems: 16 + minItems: 1 + type: array + x-kubernetes-list-map-keys: + - name + x-kubernetes-list-type: map + args: + items: + type: string + maxItems: 64 + minItems: 1 + type: array + command: + items: + type: string + maxItems: 64 + minItems: 1 + type: array + containerName: + type: string + env: + items: + properties: + name: + type: string + value: + type: string + valueFrom: properties: - claims: - items: - properties: - name: - type: string - request: - type: string - required: - - name - type: object - type: array - x-kubernetes-list-map-keys: - - name - x-kubernetes-list-type: map - limits: - additionalProperties: - anyOf: - - type: integer - - type: string - pattern: ^(\+|-)?(([0-9]+(\.[0-9]*)?)|(\.[0-9]+))(([KMGTPE]i)|[numkMGTPE]|([eE](\+|-)?(([0-9]+(\.[0-9]*)?)|(\.[0-9]+))))?$ - x-kubernetes-int-or-string: true + configMapKeyRef: + properties: + key: + type: string + name: + default: "" + type: string + optional: + type: boolean + required: + - key type: object - requests: - additionalProperties: - anyOf: - - type: integer - - type: string - pattern: ^(\+|-)?(([0-9]+(\.[0-9]*)?)|(\.[0-9]+))(([KMGTPE]i)|[numkMGTPE]|([eE](\+|-)?(([0-9]+(\.[0-9]*)?)|(\.[0-9]+))))?$ - x-kubernetes-int-or-string: true + x-kubernetes-map-type: atomic + fieldRef: + properties: + apiVersion: + type: string + fieldPath: + type: string + required: + - fieldPath + type: object + x-kubernetes-map-type: atomic + fileKeyRef: + properties: + key: + type: string + optional: + default: false + type: boolean + path: + type: string + volumeName: + type: string + required: + - key + - path + - volumeName + type: object + x-kubernetes-map-type: atomic + resourceFieldRef: + properties: + containerName: + type: string + divisor: + anyOf: + - type: integer + - type: string + pattern: ^(\+|-)?(([0-9]+(\.[0-9]*)?)|(\.[0-9]+))(([KMGTPE]i)|[numkMGTPE]|([eE](\+|-)?(([0-9]+(\.[0-9]*)?)|(\.[0-9]+))))?$ + x-kubernetes-int-or-string: true + resource: + type: string + required: + - resource + type: object + x-kubernetes-map-type: atomic + secretKeyRef: + properties: + key: + type: string + name: + default: "" + type: string + optional: + type: boolean + required: + - key type: object + x-kubernetes-map-type: atomic type: object - timeoutSeconds: - format: int64 - maximum: 3600 - minimum: 1 - type: integer + required: + - name type: object - maxProperties: 16 - minProperties: 1 + maxItems: 128 + type: array + resources: + properties: + claims: + items: + properties: + name: + type: string + request: + type: string + required: + - name + type: object + type: array + x-kubernetes-list-map-keys: + - name + x-kubernetes-list-type: map + limits: + additionalProperties: + anyOf: + - type: integer + - type: string + pattern: ^(\+|-)?(([0-9]+(\.[0-9]*)?)|(\.[0-9]+))(([KMGTPE]i)|[numkMGTPE]|([eE](\+|-)?(([0-9]+(\.[0-9]*)?)|(\.[0-9]+))))?$ + x-kubernetes-int-or-string: true + type: object + requests: + additionalProperties: + anyOf: + - type: integer + - type: string + pattern: ^(\+|-)?(([0-9]+(\.[0-9]*)?)|(\.[0-9]+))(([KMGTPE]i)|[numkMGTPE]|([eE](\+|-)?(([0-9]+(\.[0-9]*)?)|(\.[0-9]+))))?$ + x-kubernetes-int-or-string: true + type: object type: object + timeoutSeconds: + format: int64 + maximum: 3600 + minimum: 1 + type: integer required: - actions type: object diff --git a/pkg/wandb/manifest/manifest.go b/pkg/wandb/manifest/manifest.go index f8070905..ad7649ab 100644 --- a/pkg/wandb/manifest/manifest.go +++ b/pkg/wandb/manifest/manifest.go @@ -178,20 +178,23 @@ type Application struct { Triage *ApplicationTriage `yaml:"triage,omitempty"` } -// ApplicationTriage declares diagnostic actions available for an application. -// The operator copies these compact overrides onto the generated Application -// CR, which remains the source of the runtime pod configuration. +// ApplicationTriage declares a shared diagnostic runner and the actions +// available for an application. The generated Application CR remains the +// source of runtime pod configuration. type ApplicationTriage struct { - Actions map[string]TriageAction `yaml:"actions"` -} - -type TriageAction struct { ContainerName string `yaml:"containerName,omitempty"` Command []string `yaml:"command,omitempty"` Args []string `yaml:"args,omitempty"` Env []corev1.EnvVar `yaml:"env,omitempty"` Resources *corev1.ResourceRequirements `yaml:"resources,omitempty"` TimeoutSeconds int64 `yaml:"timeoutSeconds,omitempty"` + Actions []TriageAction `yaml:"actions"` +} + +type TriageAction struct { + Name string `yaml:"name"` + Description string `yaml:"description,omitempty"` + Args []string `yaml:"args,omitempty"` } type AppIngressSpec struct { From 418f4fd15b691733352dd245910b31a22e86d341 Mon Sep 17 00:00:00 2001 From: Aravind Warrier Date: Wed, 5 Aug 2026 12:56:58 -0400 Subject: [PATCH 4/4] refactor(controller): make default triage action implicit --- api/v2/application_types.go | 20 +++++++++++-------- internal/controller/triagerun_controller.go | 5 ++++- .../triagerun_controller_unit_test.go | 15 +++++++++----- 3 files changed, 26 insertions(+), 14 deletions(-) diff --git a/api/v2/application_types.go b/api/v2/application_types.go index 731da7e4..8ad7e5b3 100644 --- a/api/v2/application_types.go +++ b/api/v2/application_types.go @@ -78,8 +78,9 @@ type ApplicationSpec struct { } // ApplicationTriageSpec contains the shared diagnostic runner and the actions -// exposed by an Application. TriageRun selects actions by name; the controller -// appends each selected name to the resolved runner arguments. +// exposed by an Application. TriageRun selects actions by name. The default +// action runs the declared runner unchanged; named actions receive an explicit +// --action selector. type ApplicationTriageSpec struct { // ContainerName selects a container from the Application pod template. It // may be omitted when the Application has exactly one container. @@ -92,8 +93,9 @@ type ApplicationTriageSpec struct { // +optional Command []string `json:"command,omitempty"` - // Args replaces the selected container's arguments when non-empty. The - // controller appends the selected action name and then that action's Args. + // Args replaces the selected container's arguments when non-empty. For a + // named action, the controller appends "--action", the selected action name, + // and then that action's Args. The default action only appends its Args. // +kubebuilder:validation:MinItems=1 // +kubebuilder:validation:MaxItems=64 // +optional @@ -130,8 +132,9 @@ type ApplicationTriageSpec struct { // runner. Execution identity and resource settings remain on the parent // ApplicationTriageSpec so every action uses the same bounded runtime. type TriageActionSpec struct { - // Name is the stable identifier selected by TriageRun and passed to the - // diagnostic runner as its next argument. + // Name is the stable identifier selected by TriageRun. "default" invokes the + // shared runner without an action selector; every other name is passed as + // "--action ". Name TriageActionName `json:"name"` // Description is human-readable help shown by clients such as Watchtower. @@ -139,8 +142,9 @@ type TriageActionSpec struct { // +optional Description string `json:"description,omitempty"` - // Args are appended after the action name when starting the diagnostic - // runner. They are suitable for action-specific flags, not executable paths. + // Args are appended after the optional action selector when starting the + // diagnostic runner. They are suitable for action-specific flags, not + // executable paths. // +kubebuilder:validation:MaxItems=32 // +optional Args []string `json:"args,omitempty"` diff --git a/internal/controller/triagerun_controller.go b/internal/controller/triagerun_controller.go index 99c97aa9..b1e09de4 100644 --- a/internal/controller/triagerun_controller.go +++ b/internal/controller/triagerun_controller.go @@ -55,6 +55,7 @@ const ( triageRunLabel = "apps.wandb.com/triage-run" triageApplicationLabel = "apps.wandb.com/triage-application" triageActionAnnotation = "apps.wandb.com/triage-action" + triageActionFlag = "--action" ) // TriagePodLogReader reads the structured output from a completed triage pod. @@ -396,7 +397,9 @@ func buildTriageContainer( if len(runner.Args) > 0 { container.Args = append([]string(nil), runner.Args...) } - container.Args = append(container.Args, string(action.Name)) + if action.Name != defaultTriageAction { + container.Args = append(container.Args, triageActionFlag, string(action.Name)) + } container.Args = append(container.Args, action.Args...) if runner.Resources != nil { container.Resources = *runner.Resources.DeepCopy() diff --git a/internal/controller/triagerun_controller_unit_test.go b/internal/controller/triagerun_controller_unit_test.go index 07507e52..c3b68719 100644 --- a/internal/controller/triagerun_controller_unit_test.go +++ b/internal/controller/triagerun_controller_unit_test.go @@ -3,6 +3,7 @@ package controller import ( "context" "fmt" + "slices" "testing" wandbv2 "github.com/wandb/operator/api/v2" @@ -82,9 +83,8 @@ func TestTriageRunCreatesBoundedJobFromApplication(t *testing.T) { ) >= 0 { t.Fatal("triage memory request must be smaller than the parent application request") } - if len(container.Args) != 4 || container.Args[0] != "python" || - container.Args[3] != defaultTriageAction { - t.Fatalf("args = %#v, want triage action args", container.Args) + if got, want := container.Args, []string{"python", "-m", "weave_triage"}; !slices.Equal(got, want) { + t.Fatalf("args = %#v, want default runner args %#v", got, want) } if got := envValue(container.Env, "PYTHONPATH"); got != "/weave/src" { t.Fatalf("PYTHONPATH = %q, want action override", got) @@ -137,6 +137,7 @@ func TestTriageRunCreatesOneJobPerSelectedAction(t *testing.T) { wandbv2.TriageActionSpec{ Name: "deep", Description: "Run deeper diagnostics", + Args: []string{"--verbose"}, }) fakeClient := fake.NewClientBuilder(). WithScheme(testScheme). @@ -168,8 +169,12 @@ func TestTriageRunCreatesOneJobPerSelectedAction(t *testing.T) { job.Name, job.Annotations[triageActionAnnotation], action) } args := job.Spec.Template.Spec.Containers[0].Args - if len(args) != 4 || args[3] != action { - t.Fatalf("Job %q args = %#v, want selected action %q appended", job.Name, args, action) + want := []string{"python", "-m", "weave_triage"} + if action != defaultTriageAction { + want = append(want, triageActionFlag, action, "--verbose") + } + if !slices.Equal(args, want) { + t.Fatalf("Job %q args = %#v, want %#v", job.Name, args, want) } }