From c7fff015147cc5c04664ab9b37dbaabb5f4a45e6 Mon Sep 17 00:00:00 2001 From: Aravind Warrier Date: Tue, 28 Jul 2026 16:23:12 -0400 Subject: [PATCH 1/6] feat(controller): Reconcile TriageRun jobs --- api/v2/application_types.go | 56 ++ 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 + 13 files changed, 1571 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/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 545c49bf..e566a56b 100644 --- a/internal/controller/reconciler/reconcile_v2.go +++ b/internal/controller/reconciler/reconcile_v2.go @@ -565,6 +565,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 ce1cd96a..8e44c059 100644 --- a/pkg/wandb/manifest/manifest.go +++ b/pkg/wandb/manifest/manifest.go @@ -179,6 +179,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 f2fcc5263e1c8db36e4e0a64ca2eb3553557ed63 Mon Sep 17 00:00:00 2001 From: Aravind Warrier Date: Tue, 4 Aug 2026 11:43:00 -0400 Subject: [PATCH 2/6] 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 0e2dbf1aa3b132ec6e583482cfd2dafa23487090 Mon Sep 17 00:00:00 2001 From: Aravind Warrier Date: Tue, 4 Aug 2026 12:29:14 -0400 Subject: [PATCH 3/6] refactor(controller): resolve structured triage actions --- api/v2/application_types.go | 46 ++- api/v2/zz_generated.deepcopy.go | 335 ++++++++++++++++++ .../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, 817 insertions(+), 383 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 a6c4d885..801e7ffb 100644 --- a/api/v2/zz_generated.deepcopy.go +++ b/api/v2/zz_generated.deepcopy.go @@ -483,6 +483,53 @@ func (in *ApplicationStatus) DeepCopy() *ApplicationStatus { } // DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. +<<<<<<< HEAD +======= +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([]TriageActionSpec, len(*in)) + for i := range *in { + (*in)[i].DeepCopyInto(&(*out)[i]) + } + } +} + +// 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. +>>>>>>> a205da8 (refactor(controller): resolve structured triage actions) func (in *CertManagerConfig) DeepCopyInto(out *CertManagerConfig) { *out = *in } @@ -1736,6 +1783,294 @@ func (in *TelemetryInfraStatus) DeepCopy() *TelemetryInfraStatus { } // DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. +<<<<<<< HEAD +======= +func (in *TriageActionReference) DeepCopyInto(out *TriageActionReference) { + *out = *in +} + +// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new TriageActionReference. +func (in *TriageActionReference) DeepCopy() *TriageActionReference { + if in == nil { + return nil + } + out := new(TriageActionReference) + in.DeepCopyInto(out) + 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.Args != nil { + in, out := &in.Args, &out.Args + *out = make([]string, len(*in)) + copy(*out, *in) + } +} + +// 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 + if in.JobRef != nil { + in, out := &in.JobRef, &out.JobRef + *out = new(v1.LocalObjectReference) + **out = **in + } + if in.ResolvedExecution != nil { + in, out := &in.ResolvedExecution, &out.ResolvedExecution + *out = new(TriageResolvedExecution) + (*in).DeepCopyInto(*out) + } + if in.StartedAt != nil { + in, out := &in.StartedAt, &out.StartedAt + *out = (*in).DeepCopy() + } + if in.CompletedAt != nil { + in, out := &in.CompletedAt, &out.CompletedAt + *out = (*in).DeepCopy() + } + if in.Summary != nil { + in, out := &in.Summary, &out.Summary + *out = new(TriageRunSummary) + **out = **in + } + if in.Results != nil { + in, out := &in.Results, &out.Results + *out = make([]TriageCheckResult, len(*in)) + for i := range *in { + (*in)[i].DeepCopyInto(&(*out)[i]) + } + } +} + +// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new TriageActionStatus. +func (in *TriageActionStatus) DeepCopy() *TriageActionStatus { + if in == nil { + return nil + } + out := new(TriageActionStatus) + in.DeepCopyInto(out) + return out +} + +// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. +func (in *TriageApplicationReference) DeepCopyInto(out *TriageApplicationReference) { + *out = *in +} + +// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new TriageApplicationReference. +func (in *TriageApplicationReference) DeepCopy() *TriageApplicationReference { + if in == nil { + return nil + } + out := new(TriageApplicationReference) + in.DeepCopyInto(out) + return out +} + +// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. +func (in *TriageCheckResult) DeepCopyInto(out *TriageCheckResult) { + *out = *in + if in.Evidence != nil { + in, out := &in.Evidence, &out.Evidence + *out = new(apiextensionsv1.JSON) + (*in).DeepCopyInto(*out) + } + if in.StartedAt != nil { + in, out := &in.StartedAt, &out.StartedAt + *out = (*in).DeepCopy() + } + if in.EndedAt != nil { + in, out := &in.EndedAt, &out.EndedAt + *out = (*in).DeepCopy() + } +} + +// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new TriageCheckResult. +func (in *TriageCheckResult) DeepCopy() *TriageCheckResult { + if in == nil { + return nil + } + out := new(TriageCheckResult) + in.DeepCopyInto(out) + return out +} + +// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. +func (in *TriageResolvedExecution) DeepCopyInto(out *TriageResolvedExecution) { + *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) + } +} + +// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new TriageResolvedExecution. +func (in *TriageResolvedExecution) DeepCopy() *TriageResolvedExecution { + if in == nil { + return nil + } + out := new(TriageResolvedExecution) + in.DeepCopyInto(out) + return out +} + +// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. +func (in *TriageRun) DeepCopyInto(out *TriageRun) { + *out = *in + out.TypeMeta = in.TypeMeta + in.ObjectMeta.DeepCopyInto(&out.ObjectMeta) + in.Spec.DeepCopyInto(&out.Spec) + in.Status.DeepCopyInto(&out.Status) +} + +// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new TriageRun. +func (in *TriageRun) DeepCopy() *TriageRun { + if in == nil { + return nil + } + out := new(TriageRun) + in.DeepCopyInto(out) + return out +} + +// DeepCopyObject is an autogenerated deepcopy function, copying the receiver, creating a new runtime.Object. +func (in *TriageRun) DeepCopyObject() runtime.Object { + if c := in.DeepCopy(); c != nil { + return c + } + return nil +} + +// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. +func (in *TriageRunList) DeepCopyInto(out *TriageRunList) { + *out = *in + out.TypeMeta = in.TypeMeta + in.ListMeta.DeepCopyInto(&out.ListMeta) + if in.Items != nil { + in, out := &in.Items, &out.Items + *out = make([]TriageRun, len(*in)) + for i := range *in { + (*in)[i].DeepCopyInto(&(*out)[i]) + } + } +} + +// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new TriageRunList. +func (in *TriageRunList) DeepCopy() *TriageRunList { + if in == nil { + return nil + } + out := new(TriageRunList) + in.DeepCopyInto(out) + return out +} + +// DeepCopyObject is an autogenerated deepcopy function, copying the receiver, creating a new runtime.Object. +func (in *TriageRunList) DeepCopyObject() runtime.Object { + if c := in.DeepCopy(); c != nil { + return c + } + return nil +} + +// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. +func (in *TriageRunSpec) DeepCopyInto(out *TriageRunSpec) { + *out = *in + out.ApplicationRef = in.ApplicationRef + if in.Actions != nil { + in, out := &in.Actions, &out.Actions + *out = make([]TriageActionReference, len(*in)) + copy(*out, *in) + } +} + +// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new TriageRunSpec. +func (in *TriageRunSpec) DeepCopy() *TriageRunSpec { + if in == nil { + return nil + } + out := new(TriageRunSpec) + in.DeepCopyInto(out) + return out +} + +// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. +func (in *TriageRunStatus) DeepCopyInto(out *TriageRunStatus) { + *out = *in + if in.StartedAt != nil { + in, out := &in.StartedAt, &out.StartedAt + *out = (*in).DeepCopy() + } + if in.CompletedAt != nil { + in, out := &in.CompletedAt, &out.CompletedAt + *out = (*in).DeepCopy() + } + if in.Summary != nil { + in, out := &in.Summary, &out.Summary + *out = new(TriageRunSummary) + **out = **in + } + if in.ActionStatuses != nil { + in, out := &in.ActionStatuses, &out.ActionStatuses + *out = make([]TriageActionStatus, len(*in)) + for i := range *in { + (*in)[i].DeepCopyInto(&(*out)[i]) + } + } + if in.Conditions != nil { + in, out := &in.Conditions, &out.Conditions + *out = make([]metav1.Condition, len(*in)) + for i := range *in { + (*in)[i].DeepCopyInto(&(*out)[i]) + } + } +} + +// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new TriageRunStatus. +func (in *TriageRunStatus) DeepCopy() *TriageRunStatus { + if in == nil { + return nil + } + out := new(TriageRunStatus) + in.DeepCopyInto(out) + return out +} + +// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. +func (in *TriageRunSummary) DeepCopyInto(out *TriageRunSummary) { + *out = *in +} + +// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new TriageRunSummary. +func (in *TriageRunSummary) DeepCopy() *TriageRunSummary { + if in == nil { + return nil + } + out := new(TriageRunSummary) + in.DeepCopyInto(out) + return out +} + +// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. +>>>>>>> a205da8 (refactor(controller): resolve structured triage actions) func (in *WBInfraStatus) DeepCopyInto(out *WBInfraStatus) { *out = *in if in.Conditions != nil { 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 8e44c059..042ff8cb 100644 --- a/pkg/wandb/manifest/manifest.go +++ b/pkg/wandb/manifest/manifest.go @@ -182,20 +182,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 4d198ac089f4a83d7d1b12aaaa1d7d7fd5445ade Mon Sep 17 00:00:00 2001 From: Aravind Warrier Date: Wed, 5 Aug 2026 12:56:58 -0400 Subject: [PATCH 4/6] 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) } } From c7c9745c4e34a0675ebecaccda388a330b1a77f6 Mon Sep 17 00:00:00 2001 From: Aravind Warrier Date: Wed, 12 Aug 2026 16:13:25 -0400 Subject: [PATCH 5/6] refactor(controller): pass triage action through environment --- api/v2/application_types.go | 20 ++++++++----------- internal/controller/triagerun_controller.go | 9 +++++---- .../triagerun_controller_unit_test.go | 16 ++++++++++----- 3 files changed, 24 insertions(+), 21 deletions(-) diff --git a/api/v2/application_types.go b/api/v2/application_types.go index 8ad7e5b3..3fcb6fc3 100644 --- a/api/v2/application_types.go +++ b/api/v2/application_types.go @@ -78,9 +78,8 @@ type ApplicationSpec struct { } // ApplicationTriageSpec contains the shared diagnostic runner and the actions -// exposed by an Application. TriageRun selects actions by name. The default -// action runs the declared runner unchanged; named actions receive an explicit -// --action selector. +// exposed by an Application. TriageRun selects actions by name. The controller +// exposes the selected action to the runner through WANDB_TRIAGE_ACTION. type ApplicationTriageSpec struct { // ContainerName selects a container from the Application pod template. It // may be omitted when the Application has exactly one container. @@ -93,9 +92,8 @@ type ApplicationTriageSpec struct { // +optional Command []string `json:"command,omitempty"` - // 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. + // Args replaces the selected container's arguments when non-empty. The + // selected action's Args are appended when starting the diagnostic runner. // +kubebuilder:validation:MinItems=1 // +kubebuilder:validation:MaxItems=64 // +optional @@ -132,9 +130,8 @@ 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. "default" invokes the - // shared runner without an action selector; every other name is passed as - // "--action ". + // Name is the stable identifier selected by TriageRun and passed to the + // shared runner through WANDB_TRIAGE_ACTION. Name TriageActionName `json:"name"` // Description is human-readable help shown by clients such as Watchtower. @@ -142,9 +139,8 @@ type TriageActionSpec struct { // +optional Description string `json:"description,omitempty"` - // Args are appended after the optional action selector when starting the - // diagnostic runner. They are suitable for action-specific flags, not - // executable paths. + // Args are appended 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 b1e09de4..539315f3 100644 --- a/internal/controller/triagerun_controller.go +++ b/internal/controller/triagerun_controller.go @@ -55,7 +55,7 @@ const ( triageRunLabel = "apps.wandb.com/triage-run" triageApplicationLabel = "apps.wandb.com/triage-application" triageActionAnnotation = "apps.wandb.com/triage-action" - triageActionFlag = "--action" + triageActionEnv = "WANDB_TRIAGE_ACTION" ) // TriagePodLogReader reads the structured output from a completed triage pod. @@ -397,10 +397,11 @@ func buildTriageContainer( if len(runner.Args) > 0 { container.Args = append([]string(nil), runner.Args...) } - if action.Name != defaultTriageAction { - container.Args = append(container.Args, triageActionFlag, string(action.Name)) - } container.Args = append(container.Args, action.Args...) + container.Env = mergeTriageEnv(container.Env, []corev1.EnvVar{{ + Name: triageActionEnv, + Value: string(action.Name), + }}) 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 c3b68719..bf4a6c31 100644 --- a/internal/controller/triagerun_controller_unit_test.go +++ b/internal/controller/triagerun_controller_unit_test.go @@ -92,6 +92,9 @@ func TestTriageRunCreatesBoundedJobFromApplication(t *testing.T) { if got := envValue(container.Env, "DATABASE_URL"); got != "mysql://wandb" { t.Fatalf("DATABASE_URL = %q, want inherited env", got) } + if got := envValue(container.Env, triageActionEnv); got != defaultTriageAction { + t.Fatalf("%s = %q, want %q", triageActionEnv, got, defaultTriageAction) + } if len(container.Ports) != 0 || container.ReadinessProbe != nil || container.LivenessProbe != nil { t.Fatal("triage container must not inherit serving ports or probes") } @@ -171,11 +174,14 @@ func TestTriageRunCreatesOneJobPerSelectedAction(t *testing.T) { args := job.Spec.Template.Spec.Containers[0].Args want := []string{"python", "-m", "weave_triage"} if action != defaultTriageAction { - want = append(want, triageActionFlag, action, "--verbose") + want = append(want, "--verbose") } if !slices.Equal(args, want) { t.Fatalf("Job %q args = %#v, want %#v", job.Name, args, want) } + if got := envValue(job.Spec.Template.Spec.Containers[0].Env, triageActionEnv); got != action { + t.Fatalf("Job %q %s = %q, want %q", job.Name, triageActionEnv, got, action) + } } var updatedRun wandbv2.TriageRun @@ -442,10 +448,10 @@ func testTriageApplication() *wandbv2.Application { Triage: &wandbv2.ApplicationTriageSpec{ ContainerName: "weave-trace", Args: []string{"python", "-m", "weave_triage"}, - Env: []corev1.EnvVar{{ - Name: "PYTHONPATH", - Value: "/weave/src", - }}, + Env: []corev1.EnvVar{ + {Name: "PYTHONPATH", Value: "/weave/src"}, + {Name: triageActionEnv, Value: "manifest-value-must-not-win"}, + }, Actions: []wandbv2.TriageActionSpec{{ Name: defaultTriageAction, Description: "Run all diagnostics", From 785a289a0730cb8184a69542d898bb32710bb63b Mon Sep 17 00:00:00 2001 From: Aravind Warrier Date: Thu, 13 Aug 2026 16:50:07 -0400 Subject: [PATCH 6/6] refactor(controller): reconcile singular ActionRun jobs --- api/v2/application_types.go | 19 +- api/v2/zz_generated.deepcopy.go | 322 +------ cmd/manager/main.go | 8 +- config/rbac/role.yaml | 30 +- .../templates/wandb-operator-wandb-role.yaml | 4 +- internal/controller/actionrun_controller.go | 765 +++++++++++++++ .../actionrun_controller_unit_test.go | 434 +++++++++ internal/controller/reconciler/triage.go | 6 +- internal/controller/triagerun_controller.go | 889 ------------------ .../triagerun_controller_unit_test.go | 475 ---------- 10 files changed, 1261 insertions(+), 1691 deletions(-) create mode 100644 internal/controller/actionrun_controller.go create mode 100644 internal/controller/actionrun_controller_unit_test.go delete mode 100644 internal/controller/triagerun_controller.go delete mode 100644 internal/controller/triagerun_controller_unit_test.go diff --git a/api/v2/application_types.go b/api/v2/application_types.go index 3fcb6fc3..46decaca 100644 --- a/api/v2/application_types.go +++ b/api/v2/application_types.go @@ -68,7 +68,7 @@ type ApplicationSpec struct { CronJobs []batchv1.CronJob `json:"cronJobs,omitempty"` // Triage declares the bounded diagnostic actions that may be requested for - // this application through TriageRun resources. + // this application through ActionRun resources whose type is triage. // +optional Triage *ApplicationTriageSpec `json:"triage,omitempty"` @@ -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 -// exposes the selected action to the runner through WANDB_TRIAGE_ACTION. +// exposed by an Application. An ActionRun selects one action by name. The +// controller exposes its type and name through WANDB_ACTION_TYPE and +// WANDB_ACTION_NAME. type ApplicationTriageSpec struct { // ContainerName selects a container from the Application pod template. It // may be omitted when the Application has exactly one container. @@ -123,16 +124,16 @@ type ApplicationTriageSpec struct { // +kubebuilder:validation:MaxItems=16 // +listType=map // +listMapKey=name - Actions []TriageActionSpec `json:"actions"` + Actions []ApplicationActionSpec `json:"actions"` } -// TriageActionSpec describes one action exposed by the shared diagnostic +// ApplicationActionSpec describes one action exposed by a shared application // 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 - // shared runner through WANDB_TRIAGE_ACTION. - Name TriageActionName `json:"name"` +type ApplicationActionSpec struct { + // Name is the stable identifier selected by ActionRun and passed to the + // shared runner through WANDB_ACTION_NAME. + Name ActionName `json:"name"` // Description is human-readable help shown by clients such as Watchtower. // +kubebuilder:validation:MaxLength=512 diff --git a/api/v2/zz_generated.deepcopy.go b/api/v2/zz_generated.deepcopy.go index 801e7ffb..fc72cfd5 100644 --- a/api/v2/zz_generated.deepcopy.go +++ b/api/v2/zz_generated.deepcopy.go @@ -273,6 +273,26 @@ func (in *Application) DeepCopyObject() runtime.Object { return nil } +// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. +func (in *ApplicationActionSpec) DeepCopyInto(out *ApplicationActionSpec) { + *out = *in + if in.Args != nil { + in, out := &in.Args, &out.Args + *out = make([]string, len(*in)) + copy(*out, *in) + } +} + +// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new ApplicationActionSpec. +func (in *ApplicationActionSpec) DeepCopy() *ApplicationActionSpec { + if in == nil { + return nil + } + out := new(ApplicationActionSpec) + in.DeepCopyInto(out) + return out +} + // DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. func (in *ApplicationAutoscalingOverride) DeepCopyInto(out *ApplicationAutoscalingOverride) { *out = *in @@ -401,6 +421,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) @@ -483,8 +508,6 @@ func (in *ApplicationStatus) DeepCopy() *ApplicationStatus { } // DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. -<<<<<<< HEAD -======= func (in *ApplicationTriageSpec) DeepCopyInto(out *ApplicationTriageSpec) { *out = *in if in.Command != nil { @@ -499,19 +522,19 @@ func (in *ApplicationTriageSpec) DeepCopyInto(out *ApplicationTriageSpec) { } if in.Env != nil { in, out := &in.Env, &out.Env - *out = make([]v1.EnvVar, len(*in)) + *out = make([]corev1.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) + *out = new(corev1.ResourceRequirements) (*in).DeepCopyInto(*out) } if in.Actions != nil { in, out := &in.Actions, &out.Actions - *out = make([]TriageActionSpec, len(*in)) + *out = make([]ApplicationActionSpec, len(*in)) for i := range *in { (*in)[i].DeepCopyInto(&(*out)[i]) } @@ -529,7 +552,6 @@ func (in *ApplicationTriageSpec) DeepCopy() *ApplicationTriageSpec { } // DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. ->>>>>>> a205da8 (refactor(controller): resolve structured triage actions) func (in *CertManagerConfig) DeepCopyInto(out *CertManagerConfig) { *out = *in } @@ -1783,294 +1805,6 @@ func (in *TelemetryInfraStatus) DeepCopy() *TelemetryInfraStatus { } // DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. -<<<<<<< HEAD -======= -func (in *TriageActionReference) DeepCopyInto(out *TriageActionReference) { - *out = *in -} - -// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new TriageActionReference. -func (in *TriageActionReference) DeepCopy() *TriageActionReference { - if in == nil { - return nil - } - out := new(TriageActionReference) - in.DeepCopyInto(out) - 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.Args != nil { - in, out := &in.Args, &out.Args - *out = make([]string, len(*in)) - copy(*out, *in) - } -} - -// 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 - if in.JobRef != nil { - in, out := &in.JobRef, &out.JobRef - *out = new(v1.LocalObjectReference) - **out = **in - } - if in.ResolvedExecution != nil { - in, out := &in.ResolvedExecution, &out.ResolvedExecution - *out = new(TriageResolvedExecution) - (*in).DeepCopyInto(*out) - } - if in.StartedAt != nil { - in, out := &in.StartedAt, &out.StartedAt - *out = (*in).DeepCopy() - } - if in.CompletedAt != nil { - in, out := &in.CompletedAt, &out.CompletedAt - *out = (*in).DeepCopy() - } - if in.Summary != nil { - in, out := &in.Summary, &out.Summary - *out = new(TriageRunSummary) - **out = **in - } - if in.Results != nil { - in, out := &in.Results, &out.Results - *out = make([]TriageCheckResult, len(*in)) - for i := range *in { - (*in)[i].DeepCopyInto(&(*out)[i]) - } - } -} - -// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new TriageActionStatus. -func (in *TriageActionStatus) DeepCopy() *TriageActionStatus { - if in == nil { - return nil - } - out := new(TriageActionStatus) - in.DeepCopyInto(out) - return out -} - -// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. -func (in *TriageApplicationReference) DeepCopyInto(out *TriageApplicationReference) { - *out = *in -} - -// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new TriageApplicationReference. -func (in *TriageApplicationReference) DeepCopy() *TriageApplicationReference { - if in == nil { - return nil - } - out := new(TriageApplicationReference) - in.DeepCopyInto(out) - return out -} - -// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. -func (in *TriageCheckResult) DeepCopyInto(out *TriageCheckResult) { - *out = *in - if in.Evidence != nil { - in, out := &in.Evidence, &out.Evidence - *out = new(apiextensionsv1.JSON) - (*in).DeepCopyInto(*out) - } - if in.StartedAt != nil { - in, out := &in.StartedAt, &out.StartedAt - *out = (*in).DeepCopy() - } - if in.EndedAt != nil { - in, out := &in.EndedAt, &out.EndedAt - *out = (*in).DeepCopy() - } -} - -// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new TriageCheckResult. -func (in *TriageCheckResult) DeepCopy() *TriageCheckResult { - if in == nil { - return nil - } - out := new(TriageCheckResult) - in.DeepCopyInto(out) - return out -} - -// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. -func (in *TriageResolvedExecution) DeepCopyInto(out *TriageResolvedExecution) { - *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) - } -} - -// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new TriageResolvedExecution. -func (in *TriageResolvedExecution) DeepCopy() *TriageResolvedExecution { - if in == nil { - return nil - } - out := new(TriageResolvedExecution) - in.DeepCopyInto(out) - return out -} - -// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. -func (in *TriageRun) DeepCopyInto(out *TriageRun) { - *out = *in - out.TypeMeta = in.TypeMeta - in.ObjectMeta.DeepCopyInto(&out.ObjectMeta) - in.Spec.DeepCopyInto(&out.Spec) - in.Status.DeepCopyInto(&out.Status) -} - -// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new TriageRun. -func (in *TriageRun) DeepCopy() *TriageRun { - if in == nil { - return nil - } - out := new(TriageRun) - in.DeepCopyInto(out) - return out -} - -// DeepCopyObject is an autogenerated deepcopy function, copying the receiver, creating a new runtime.Object. -func (in *TriageRun) DeepCopyObject() runtime.Object { - if c := in.DeepCopy(); c != nil { - return c - } - return nil -} - -// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. -func (in *TriageRunList) DeepCopyInto(out *TriageRunList) { - *out = *in - out.TypeMeta = in.TypeMeta - in.ListMeta.DeepCopyInto(&out.ListMeta) - if in.Items != nil { - in, out := &in.Items, &out.Items - *out = make([]TriageRun, len(*in)) - for i := range *in { - (*in)[i].DeepCopyInto(&(*out)[i]) - } - } -} - -// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new TriageRunList. -func (in *TriageRunList) DeepCopy() *TriageRunList { - if in == nil { - return nil - } - out := new(TriageRunList) - in.DeepCopyInto(out) - return out -} - -// DeepCopyObject is an autogenerated deepcopy function, copying the receiver, creating a new runtime.Object. -func (in *TriageRunList) DeepCopyObject() runtime.Object { - if c := in.DeepCopy(); c != nil { - return c - } - return nil -} - -// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. -func (in *TriageRunSpec) DeepCopyInto(out *TriageRunSpec) { - *out = *in - out.ApplicationRef = in.ApplicationRef - if in.Actions != nil { - in, out := &in.Actions, &out.Actions - *out = make([]TriageActionReference, len(*in)) - copy(*out, *in) - } -} - -// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new TriageRunSpec. -func (in *TriageRunSpec) DeepCopy() *TriageRunSpec { - if in == nil { - return nil - } - out := new(TriageRunSpec) - in.DeepCopyInto(out) - return out -} - -// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. -func (in *TriageRunStatus) DeepCopyInto(out *TriageRunStatus) { - *out = *in - if in.StartedAt != nil { - in, out := &in.StartedAt, &out.StartedAt - *out = (*in).DeepCopy() - } - if in.CompletedAt != nil { - in, out := &in.CompletedAt, &out.CompletedAt - *out = (*in).DeepCopy() - } - if in.Summary != nil { - in, out := &in.Summary, &out.Summary - *out = new(TriageRunSummary) - **out = **in - } - if in.ActionStatuses != nil { - in, out := &in.ActionStatuses, &out.ActionStatuses - *out = make([]TriageActionStatus, len(*in)) - for i := range *in { - (*in)[i].DeepCopyInto(&(*out)[i]) - } - } - if in.Conditions != nil { - in, out := &in.Conditions, &out.Conditions - *out = make([]metav1.Condition, len(*in)) - for i := range *in { - (*in)[i].DeepCopyInto(&(*out)[i]) - } - } -} - -// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new TriageRunStatus. -func (in *TriageRunStatus) DeepCopy() *TriageRunStatus { - if in == nil { - return nil - } - out := new(TriageRunStatus) - in.DeepCopyInto(out) - return out -} - -// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. -func (in *TriageRunSummary) DeepCopyInto(out *TriageRunSummary) { - *out = *in -} - -// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new TriageRunSummary. -func (in *TriageRunSummary) DeepCopy() *TriageRunSummary { - if in == nil { - return nil - } - out := new(TriageRunSummary) - in.DeepCopyInto(out) - return out -} - -// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. ->>>>>>> a205da8 (refactor(controller): resolve structured triage actions) func (in *WBInfraStatus) DeepCopyInto(out *WBInfraStatus) { *out = *in if in.Conditions != nil { diff --git a/cmd/manager/main.go b/cmd/manager/main.go index eb91036d..ca366624 100644 --- a/cmd/manager/main.go +++ b/cmd/manager/main.go @@ -342,17 +342,17 @@ func main() { kubernetesClient, err := kubernetes.NewForConfig(mgr.GetConfig()) if err != nil { - setupLog.Error(err, "unable to create Kubernetes client", "controller", "TriageRun") + setupLog.Error(err, "unable to create Kubernetes client", "controller", "ActionRun") os.Exit(1) } - if err = (&controller.TriageRunReconciler{ + if err = (&controller.ActionRunReconciler{ Client: mgr.GetClient(), Scheme: mgr.GetScheme(), - PodLogs: &controller.KubernetesTriagePodLogReader{ + PodLogs: &controller.KubernetesActionPodLogReader{ CoreV1: kubernetesClient.CoreV1(), }, }).SetupWithManager(mgr); err != nil { - setupLog.Error(err, "unable to create controller", "controller", "TriageRun") + setupLog.Error(err, "unable to create controller", "controller", "ActionRun") os.Exit(1) } diff --git a/config/rbac/role.yaml b/config/rbac/role.yaml index 21e7862f..93cf69fa 100644 --- a/config/rbac/role.yaml +++ b/config/rbac/role.yaml @@ -74,28 +74,16 @@ rules: - apiGroups: - apps.wandb.com resources: - - applications - - weightsandbiases + - actionruns verbs: - - create - - delete - get - list - - patch - - update - watch - apiGroups: - apps.wandb.com resources: - - applications/finalizers - - weightsandbiases/finalizers - verbs: - - update -- apiGroups: - - apps.wandb.com - resources: + - actionruns/status - applications/status - - triageruns/status - weightsandbiases/status verbs: - get @@ -104,11 +92,23 @@ rules: - apiGroups: - apps.wandb.com resources: - - triageruns + - applications + - weightsandbiases verbs: + - create + - delete - get - list + - patch + - update - watch +- apiGroups: + - apps.wandb.com + resources: + - applications/finalizers + - weightsandbiases/finalizers + verbs: + - update - apiGroups: - autoscaling resources: diff --git a/deploy/operator/templates/wandb-operator-wandb-role.yaml b/deploy/operator/templates/wandb-operator-wandb-role.yaml index 75a26d83..49229ca3 100644 --- a/deploy/operator/templates/wandb-operator-wandb-role.yaml +++ b/deploy/operator/templates/wandb-operator-wandb-role.yaml @@ -20,7 +20,7 @@ rules: - apiGroups: - apps.wandb.com resources: - - triageruns + - actionruns verbs: - get - list @@ -36,7 +36,7 @@ rules: - apps.wandb.com resources: - applications/status - - triageruns/status + - actionruns/status - weightsandbiases/status verbs: - get diff --git a/internal/controller/actionrun_controller.go b/internal/controller/actionrun_controller.go new file mode 100644 index 00000000..d64f87a9 --- /dev/null +++ b/internal/controller/actionrun_controller.go @@ -0,0 +1,765 @@ +/* +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 ( + defaultActionTimeoutSeconds = int64(300) + maxActionOutputBytes = int64(512 * 1024) + actionContainerName = "action" + actionConditionSucceeded = "Succeeded" + actionRunLabel = "apps.wandb.com/action-run" + actionApplicationLabel = "apps.wandb.com/action-application" + actionTypeAnnotation = "apps.wandb.com/action-type" + actionNameAnnotation = "apps.wandb.com/action-name" + actionTypeEnv = "WANDB_ACTION_TYPE" + actionNameEnv = "WANDB_ACTION_NAME" +) + +// ActionPodLogReader reads structured output from a completed action pod. It +// is an interface so controller behavior can be tested without a live API +// server's pod log subresource. +type ActionPodLogReader interface { + ReadPodLogs( + ctx context.Context, + namespace string, + podName string, + containerName string, + maxBytes int64, + ) ([]byte, error) +} + +type KubernetesActionPodLogReader struct { + CoreV1 corev1client.CoreV1Interface +} + +func (r *KubernetesActionPodLogReader) 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, &actionOutputTooLargeError{maxBytes: maxBytes} + } + return output, nil +} + +// ActionRunReconciler turns one immutable ActionRun into one bounded Job and +// records the Job's structured JSONL output on the run status. +type ActionRunReconciler struct { + client.Client + Scheme *runtime.Scheme + PodLogs ActionPodLogReader +} + +// +kubebuilder:rbac:groups=apps.wandb.com,resources=actionruns,verbs=get;list;watch +// +kubebuilder:rbac:groups=apps.wandb.com,resources=actionruns/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 *ActionRunReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) { + var run wandbv2.ActionRun + if err := r.Get(ctx, req.NamespacedName, &run); err != nil { + return ctrl.Result{}, client.IgnoreNotFound(err) + } + + if isTerminalActionPhase(run.Status.Phase) { + return ctrl.Result{}, nil + } + if run.Spec.Type != wandbv2.ActionTypeTriage { + return r.failRun(ctx, &run, "UnsupportedActionType", + fmt.Sprintf("action type %q is not executable in this release", run.Spec.Type)) + } + + 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 + } + + action, err := resolveRequestedAction(&run, &application) + if err != nil { + return r.failRun(ctx, &run, "InvalidAction", err.Error()) + } + job, err := r.getOrCreateActionJob(ctx, &run, &application, action) + if err != nil { + var missing *actionJobMissingError + 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 ActionRun", job.Name)) + } + + statusBefore := run.DeepCopy().Status + requeueForOutput := r.updateRunStatus(ctx, &run, action, job) + if !apiequality.Semantic.DeepEqual(statusBefore, run.Status) { + if err := r.Status().Update(ctx, &run); err != nil { + return ctrl.Result{}, err + } + } + if requeueForOutput { + return ctrl.Result{RequeueAfter: 2 * time.Second}, nil + } + return ctrl.Result{}, nil +} + +type resolvedAction struct { + actionType wandbv2.ActionType + name wandbv2.ActionName + runner *wandbv2.ApplicationTriageSpec + spec wandbv2.ApplicationActionSpec + source *corev1.Container + timeoutSeconds int64 + resolved *wandbv2.ActionResolvedExecution + jobName string +} + +func resolveRequestedAction( + run *wandbv2.ActionRun, + application *wandbv2.Application, +) (*resolvedAction, error) { + actionName := run.Spec.Action.Name + actionSpec, err := resolveTriageAction(application, actionName) + if err != nil { + return nil, err + } + runner := application.Spec.Triage + source, err := selectActionContainer(application, runner.ContainerName) + if err != nil { + return nil, err + } + timeoutSeconds := runner.TimeoutSeconds + if timeoutSeconds == 0 { + timeoutSeconds = defaultActionTimeoutSeconds + } + action := &resolvedAction{ + actionType: run.Spec.Type, + name: actionName, + runner: runner, + spec: actionSpec, + source: source, + timeoutSeconds: timeoutSeconds, + jobName: common.FitDefaultInfraName(run.Name, "-action", 63), + } + action.resolved = resolvedActionExecution(application, action) + return action, nil +} + +func (r *ActionRunReconciler) getOrCreateActionJob( + ctx context.Context, + run *wandbv2.ActionRun, + application *wandbv2.Application, + action *resolvedAction, +) (*batchv1.Job, error) { + var job batchv1.Job + err := r.Get(ctx, types.NamespacedName{Namespace: run.Namespace, Name: action.jobName}, &job) + if err != nil && !apierrors.IsNotFound(err) { + return nil, err + } + if apierrors.IsNotFound(err) { + if run.Status.JobRef != nil { + return nil, &actionJobMissingError{name: run.Status.JobRef.Name} + } + job = *buildActionJob(run, application, action) + 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 nil, err + } + if err := r.Get(ctx, types.NamespacedName{ + Namespace: run.Namespace, + Name: action.jobName, + }, &job); err != nil { + return nil, err + } + } + } + return &job, nil +} + +type actionJobMissingError struct { + name string +} + +func (e *actionJobMissingError) Error() string { + return fmt.Sprintf("Job %q disappeared before the action completed", e.name) +} + +func resolveTriageAction( + application *wandbv2.Application, + actionName wandbv2.ActionName, +) (wandbv2.ApplicationActionSpec, error) { + if application.Spec.Triage == nil { + return wandbv2.ApplicationActionSpec{}, fmt.Errorf( + "Application %q does not declare triage actions", application.Name) + } + if len(application.Spec.Triage.Command) == 0 && len(application.Spec.Triage.Args) == 0 { + return wandbv2.ApplicationActionSpec{}, fmt.Errorf( + "Application %q triage runner must override command or args", application.Name) + } + for i := range application.Spec.Triage.Actions { + action := application.Spec.Triage.Actions[i] + if action.Name == actionName { + return action, nil + } + } + return wandbv2.ApplicationActionSpec{}, fmt.Errorf( + "Application %q does not declare triage action %q", application.Name, actionName) +} + +func selectActionContainer(application *wandbv2.Application, name string) (*corev1.Container, error) { + containers := application.Spec.PodTemplate.Spec.Containers + if name == "" { + if len(containers) != 1 { + return nil, fmt.Errorf( + "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 buildActionJob( + run *wandbv2.ActionRun, + application *wandbv2.Application, + action *resolvedAction, +) *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{buildActionContainer(action.source, action.runner, action)} + + labels := map[string]string{ + actionRunLabel: common.FitDefaultInfraName(run.Name, "", 63), + actionApplicationLabel: common.FitDefaultInfraName(application.Name, "", 63), + } + annotations := map[string]string{ + actionTypeAnnotation: string(action.actionType), + actionNameAnnotation: string(action.name), + } + return &batchv1.Job{ + ObjectMeta: metav1.ObjectMeta{ + Name: action.jobName, + Namespace: run.Namespace, + Labels: labels, + Annotations: annotations, + }, + Spec: batchv1.JobSpec{ + BackoffLimit: &backoffLimit, + ActiveDeadlineSeconds: &action.timeoutSeconds, + Template: corev1.PodTemplateSpec{ + ObjectMeta: metav1.ObjectMeta{Labels: labels, Annotations: annotations}, + Spec: *podSpec, + }, + }, + } +} + +func buildActionContainer( + source *corev1.Container, + runner *wandbv2.ApplicationTriageSpec, + action *resolvedAction, +) corev1.Container { + container := corev1.Container{ + Name: actionContainerName, + 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: mergeActionEnv(source.Env, runner.Env), + Resources: defaultActionResources(), + 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(runner.Command) > 0 { + container.Command = append([]string(nil), runner.Command...) + container.Args = nil + } + if len(runner.Args) > 0 { + container.Args = append([]string(nil), runner.Args...) + } + container.Args = append(container.Args, action.spec.Args...) + container.Env = mergeActionEnv(container.Env, []corev1.EnvVar{ + {Name: actionTypeEnv, Value: string(action.actionType)}, + {Name: actionNameEnv, Value: string(action.name)}, + }) + if runner.Resources != nil { + container.Resources = *runner.Resources.DeepCopy() + } + return container +} + +func mergeActionEnv(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 defaultActionResources() 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 resolvedActionExecution( + application *wandbv2.Application, + action *resolvedAction, +) *wandbv2.ActionResolvedExecution { + container := buildActionContainer(action.source, action.runner, action) + return &wandbv2.ActionResolvedExecution{ + ApplicationGeneration: application.Generation, + ContainerName: action.source.Name, + Image: container.Image, + Command: append([]string(nil), container.Command...), + Args: append([]string(nil), container.Args...), + TimeoutSeconds: action.timeoutSeconds, + } +} + +func (r *ActionRunReconciler) updateRunStatus( + ctx context.Context, + run *wandbv2.ActionRun, + action *resolvedAction, + job *batchv1.Job, +) bool { + run.Status.Phase = wandbv2.ActionRunPhaseRunning + run.Status.ObservedGeneration = run.Generation + run.Status.JobRef = &corev1.LocalObjectReference{Name: job.Name} + run.Status.ResolvedExecution = action.resolved + run.Status.StartedAt = actionStartTime(run.Status.StartedAt, job) + run.Status.CompletedAt = nil + run.Status.Summary = nil + run.Status.Results = nil + + if jobFailed(job) { + run.Status.Phase = wandbv2.ActionRunPhaseFailed + run.Status.CompletedAt = actionCompletionTime(job) + message := jobConditionMessage(job, batchv1.JobFailed) + if message == "" { + message = fmt.Sprintf("Job %q failed", job.Name) + } + if results, err := r.collectActionResults(ctx, job); err == nil { + run.Status.Results = results + run.Status.Summary = summarizeActionResults(results) + } + setActionCondition(run, metav1.ConditionFalse, "JobFailed", message) + return false + } + if !jobComplete(job) { + setActionCondition(run, metav1.ConditionUnknown, "ActionRunning", + fmt.Sprintf("Job %q is running", job.Name)) + return false + } + + results, err := r.collectActionResults(ctx, job) + if err != nil { + var unavailable *actionOutputUnavailableError + if errors.As(err, &unavailable) { + setActionCondition(run, metav1.ConditionUnknown, "ResultsPending", unavailable.Error()) + return true + } + run.Status.Phase = wandbv2.ActionRunPhaseFailed + run.Status.CompletedAt = actionCompletionTime(job) + setActionCondition(run, metav1.ConditionFalse, "InvalidResults", err.Error()) + return false + } + + run.Status.Phase = wandbv2.ActionRunPhaseSucceeded + run.Status.CompletedAt = actionCompletionTime(job) + run.Status.Results = results + run.Status.Summary = summarizeActionResults(results) + setActionCondition(run, metav1.ConditionTrue, "ResultsCollected", + fmt.Sprintf("Collected %d action results", len(results))) + return false +} + +func setActionCondition( + run *wandbv2.ActionRun, + status metav1.ConditionStatus, + reason string, + message string, +) { + apimeta.SetStatusCondition(&run.Status.Conditions, metav1.Condition{ + Type: actionConditionSucceeded, + Status: status, + ObservedGeneration: run.Generation, + Reason: reason, + Message: message, + }) +} + +func (r *ActionRunReconciler) failRun( + ctx context.Context, + run *wandbv2.ActionRun, + reason string, + message string, +) (ctrl.Result, error) { + run.Status.Phase = wandbv2.ActionRunPhaseFailed + run.Status.ObservedGeneration = run.Generation + now := metav1.Now() + run.Status.CompletedAt = &now + setActionCondition(run, metav1.ConditionFalse, reason, message) + if err := r.Status().Update(ctx, run); err != nil { + return ctrl.Result{}, err + } + return ctrl.Result{}, nil +} + +func actionStartTime(previous *metav1.Time, job *batchv1.Job) *metav1.Time { + if job.Status.StartTime != nil { + return job.Status.StartTime.DeepCopy() + } + if previous != nil { + return previous.DeepCopy() + } + if !job.CreationTimestamp.IsZero() { + return job.CreationTimestamp.DeepCopy() + } + now := metav1.Now() + return &now +} + +func actionCompletionTime(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 isTerminalActionPhase(phase wandbv2.ActionRunPhase) bool { + return phase == wandbv2.ActionRunPhaseSucceeded || phase == wandbv2.ActionRunPhaseFailed +} + +func (r *ActionRunReconciler) collectActionResults( + ctx context.Context, + job *batchv1.Job, +) ([]wandbv2.ActionResult, 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, &actionOutputUnavailableError{ + err: fmt.Errorf("list pods for Job %q: %w", job.Name, err), + } + } + if len(pods.Items) == 0 { + return nil, &actionOutputUnavailableError{ + 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, actionContainerName, maxActionOutputBytes) + if err != nil { + var tooLarge *actionOutputTooLargeError + if errors.As(err, &tooLarge) { + return nil, err + } + return nil, &actionOutputUnavailableError{ + err: fmt.Errorf("read action output: %w", err), + } + } + results, err := parseActionJSONL(output) + if err != nil { + return nil, fmt.Errorf("parse action output: %w", err) + } + return results, nil +} + +type actionOutputUnavailableError struct { + err error +} + +func (e *actionOutputUnavailableError) Error() string { + return e.err.Error() +} + +func (e *actionOutputUnavailableError) Unwrap() error { + return e.err +} + +type actionOutputTooLargeError struct { + maxBytes int64 +} + +func (e *actionOutputTooLargeError) Error() string { + return fmt.Sprintf("action output exceeds %d bytes", e.maxBytes) +} + +type actionJSONResult 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 parseActionJSONL(output []byte) ([]wandbv2.ActionResult, error) { + if int64(len(output)) > maxActionOutputBytes { + return nil, fmt.Errorf("action output exceeds %d bytes", maxActionOutputBytes) + } + + scanner := bufio.NewScanner(bytes.NewReader(output)) + scanner.Buffer(make([]byte, 64*1024), int(maxActionOutputBytes)) + results := make([]wandbv2.ActionResult, 0) + lineNumber := 0 + for scanner.Scan() { + lineNumber++ + line := bytes.TrimSpace(scanner.Bytes()) + if len(line) == 0 { + continue + } + + var raw actionJSONResult + 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("action command emitted no results") + } + return results, nil +} + +func (r actionJSONResult) toAPIResult() (wandbv2.ActionResult, error) { + if strings.TrimSpace(r.Name) == "" { + return wandbv2.ActionResult{}, errors.New("name is required") + } + severity := wandbv2.ActionSeverity(r.Severity) + switch severity { + case wandbv2.ActionSeverityPass, + wandbv2.ActionSeverityWarn, + wandbv2.ActionSeverityFail, + wandbv2.ActionSeverityError: + default: + return wandbv2.ActionResult{}, fmt.Errorf("unsupported severity %q", r.Severity) + } + + result := wandbv2.ActionResult{ + 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.ActionResult{}, errors.New("evidence is not valid JSON") + } + result.Evidence = &apiextensionsv1.JSON{Raw: append([]byte(nil), r.Evidence...)} + } + + var err error + result.StartedAt, err = parseActionTimestamp("started_at", r.StartedAt) + if err != nil { + return wandbv2.ActionResult{}, err + } + result.EndedAt, err = parseActionTimestamp("ended_at", r.EndedAt) + if err != nil { + return wandbv2.ActionResult{}, err + } + return result, nil +} + +func parseActionTimestamp(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 summarizeActionResults(results []wandbv2.ActionResult) *wandbv2.ActionRunSummary { + summary := &wandbv2.ActionRunSummary{Total: int32(len(results))} + for _, result := range results { + switch result.Severity { + case wandbv2.ActionSeverityPass: + summary.Pass++ + case wandbv2.ActionSeverityWarn: + summary.Warn++ + case wandbv2.ActionSeverityFail: + summary.Fail++ + case wandbv2.ActionSeverityError: + summary.Error++ + } + } + switch { + case summary.Error > 0: + summary.OverallSeverity = wandbv2.ActionSeverityError + case summary.Fail > 0: + summary.OverallSeverity = wandbv2.ActionSeverityFail + case summary.Warn > 0: + summary.OverallSeverity = wandbv2.ActionSeverityWarn + default: + summary.OverallSeverity = wandbv2.ActionSeverityPass + } + return summary +} + +func (r *ActionRunReconciler) SetupWithManager(mgr ctrl.Manager) error { + return ctrl.NewControllerManagedBy(mgr). + For(&wandbv2.ActionRun{}). + Owns(&batchv1.Job{}). + Named("actionrun"). + Complete(r) +} diff --git a/internal/controller/actionrun_controller_unit_test.go b/internal/controller/actionrun_controller_unit_test.go new file mode 100644 index 00000000..44d73b4c --- /dev/null +++ b/internal/controller/actionrun_controller_unit_test.go @@ -0,0 +1,434 @@ +package controller + +import ( + "context" + "fmt" + "slices" + "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/yaml" +) + +func TestActionRunCreatesBoundedJobFromApplication(t *testing.T) { + t.Parallel() + + testScheme := newActionTestScheme(t) + run := testActionRun() + application := testActionApplication() + fakeClient := fake.NewClientBuilder(). + WithScheme(testScheme). + WithStatusSubresource(&wandbv2.ActionRun{}, &batchv1.Job{}). + WithObjects(run, application). + Build() + reconciler := &ActionRunReconciler{ + Client: fakeClient, + Scheme: testScheme, + PodLogs: &staticActionLogReader{}, + } + + if _, err := reconciler.Reconcile(context.Background(), requestFor(run)); err != nil { + t.Fatalf("reconcile ActionRun: %v", err) + } + + var job batchv1.Job + if err := fakeClient.Get(context.Background(), types.NamespacedName{ + Namespace: run.Namespace, + Name: "weave-check-action", + }, &job); err != nil { + t.Fatalf("get action 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.Name != actionContainerName || container.Image != "weave:sha256-test" { + t.Fatalf("container = %#v, want action container with inherited image", container) + } + 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 container.Resources.Requests.Memory().Cmp( + *application.Spec.PodTemplate.Spec.Containers[0].Resources.Requests.Memory(), + ) >= 0 { + t.Fatal("action memory request must be smaller than the parent application request") + } + 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) + } + if got := envValue(container.Env, "DATABASE_URL"); got != "mysql://wandb" { + t.Fatalf("DATABASE_URL = %q, want inherited env", got) + } + if got := envValue(container.Env, actionTypeEnv); got != string(wandbv2.ActionTypeTriage) { + t.Fatalf("%s = %q, want %q", actionTypeEnv, got, wandbv2.ActionTypeTriage) + } + if got := envValue(container.Env, actionNameEnv); got != "default" { + t.Fatalf("%s = %q, want default", actionNameEnv, got) + } + if len(container.Ports) != 0 || container.ReadinessProbe != nil || container.LivenessProbe != nil { + t.Fatal("action container must not inherit serving ports or probes") + } + if len(container.VolumeMounts) != 1 || len(job.Spec.Template.Spec.Volumes) != 1 { + t.Fatal("action Job must inherit the selected container's mounts and parent pod volumes") + } + if !metav1.IsControlledBy(&job, run) { + t.Fatal("action Job is not controlled by the ActionRun") + } + + var updatedRun wandbv2.ActionRun + if err := fakeClient.Get(context.Background(), client.ObjectKeyFromObject(run), &updatedRun); err != nil { + t.Fatalf("get updated ActionRun: %v", err) + } + if updatedRun.Status.Phase != wandbv2.ActionRunPhaseRunning { + 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 TestActionRunSelectsOneNamedAction(t *testing.T) { + t.Parallel() + + testScheme := newActionTestScheme(t) + run := testActionRun() + run.Spec.Action.Name = "deep" + application := testActionApplication() + application.Spec.Triage.Actions = append(application.Spec.Triage.Actions, + wandbv2.ApplicationActionSpec{ + Name: "deep", + Description: "Run deeper diagnostics", + Args: []string{"--verbose"}, + }) + fakeClient := fake.NewClientBuilder(). + WithScheme(testScheme). + WithStatusSubresource(&wandbv2.ActionRun{}, &batchv1.Job{}). + WithObjects(run, application). + Build() + reconciler := &ActionRunReconciler{Client: fakeClient, Scheme: testScheme} + + if _, err := reconciler.Reconcile(context.Background(), requestFor(run)); err != nil { + t.Fatalf("reconcile named ActionRun: %v", err) + } + var jobs batchv1.JobList + if err := fakeClient.List(context.Background(), &jobs, client.InNamespace(run.Namespace)); err != nil { + t.Fatalf("list Jobs: %v", err) + } + if len(jobs.Items) != 1 { + t.Fatalf("Jobs = %d, want exactly one", len(jobs.Items)) + } + job := jobs.Items[0] + if job.Annotations[actionTypeAnnotation] != "triage" || + job.Annotations[actionNameAnnotation] != "deep" { + t.Fatalf("Job annotations = %#v", job.Annotations) + } + container := job.Spec.Template.Spec.Containers[0] + if got, want := container.Args, []string{"python", "-m", "weave_triage", "--verbose"}; !slices.Equal(got, want) { + t.Fatalf("args = %#v, want %#v", got, want) + } + if got := envValue(container.Env, actionNameEnv); got != "deep" { + t.Fatalf("%s = %q, want deep", actionNameEnv, got) + } +} + +func TestActionRunRejectsUnsupportedType(t *testing.T) { + t.Parallel() + + testScheme := newActionTestScheme(t) + run := testActionRun() + run.Spec.Type = wandbv2.ActionTypeMaintenance + fakeClient := fake.NewClientBuilder(). + WithScheme(testScheme). + WithStatusSubresource(&wandbv2.ActionRun{}, &batchv1.Job{}). + WithObjects(run). + Build() + reconciler := &ActionRunReconciler{Client: fakeClient, Scheme: testScheme} + + if _, err := reconciler.Reconcile(context.Background(), requestFor(run)); err != nil { + t.Fatalf("reconcile maintenance ActionRun: %v", err) + } + var updatedRun wandbv2.ActionRun + if err := fakeClient.Get(context.Background(), client.ObjectKeyFromObject(run), &updatedRun); err != nil { + t.Fatalf("get updated ActionRun: %v", err) + } + if updatedRun.Status.Phase != wandbv2.ActionRunPhaseFailed { + t.Fatalf("phase = %q, want Failed", updatedRun.Status.Phase) + } + if conditionReason(updatedRun.Status.Conditions, actionConditionSucceeded) != "UnsupportedActionType" { + t.Fatalf("conditions = %#v, want UnsupportedActionType", updatedRun.Status.Conditions) + } + var jobs batchv1.JobList + if err := fakeClient.List(context.Background(), &jobs); err != nil { + t.Fatalf("list Jobs: %v", err) + } + if len(jobs.Items) != 0 { + t.Fatalf("Jobs = %d, want none", len(jobs.Items)) + } +} + +func TestActionRunCollectsFailedCheckAsSuccessfulExecution(t *testing.T) { + t.Parallel() + + testScheme := newActionTestScheme(t) + run := testActionRun() + application := testActionApplication() + logReader := &staticActionLogReader{} + fakeClient := fake.NewClientBuilder(). + WithScheme(testScheme). + WithStatusSubresource(&wandbv2.ActionRun{}, &batchv1.Job{}). + WithObjects(run, application). + Build() + reconciler := &ActionRunReconciler{Client: fakeClient, Scheme: testScheme, PodLogs: logReader} + + if _, err := reconciler.Reconcile(context.Background(), requestFor(run)); err != nil { + t.Fatalf("create action Job: %v", err) + } + var job batchv1.Job + if err := fakeClient.Get(context.Background(), types.NamespacedName{ + Namespace: run.Namespace, + Name: "weave-check-action", + }, &job); err != nil { + t.Fatalf("get action Job: %v", 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 complete: %v", err) + } + pod := &corev1.Pod{ObjectMeta: metav1.ObjectMeta{ + Name: "weave-check-action-pod", + Namespace: run.Namespace, + Labels: map[string]string{"batch.kubernetes.io/job-name": job.Name}, + }} + if err := fakeClient.Create(context.Background(), pod); err != nil { + t.Fatalf("create Job pod: %v", err) + } + logReader.output = []byte( + `{"name":"starter-project","severity":"pass","message":"reachable"}` + "\n" + + `{"name":"starter-object","severity":"fail","evidence":{"missing":2},"remediation":"create starters"}` + "\n", + ) + + if _, err := reconciler.Reconcile(context.Background(), requestFor(run)); err != nil { + t.Fatalf("reconcile completed ActionRun: %v", err) + } + var updatedRun wandbv2.ActionRun + if err := fakeClient.Get(context.Background(), client.ObjectKeyFromObject(run), &updatedRun); err != nil { + t.Fatalf("get completed ActionRun: %v", err) + } + if updatedRun.Status.Phase != wandbv2.ActionRunPhaseSucceeded { + t.Fatalf("phase = %q, want Succeeded because the command completed", updatedRun.Status.Phase) + } + if updatedRun.Status.Summary == nil || + updatedRun.Status.Summary.OverallSeverity != wandbv2.ActionSeverityFail || + updatedRun.Status.Summary.Fail != 1 { + t.Fatalf("summary = %#v, want one failed check", updatedRun.Status.Summary) + } + if len(updatedRun.Status.Results) != 2 { + t.Fatalf("results = %#v, want 2", updatedRun.Status.Results) + } +} + +func TestParseActionJSONLRejectsNoisyOutput(t *testing.T) { + t.Parallel() + + _, err := parseActionJSONL([]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: + containerName: weave-trace + args: [python, -m, weave_triage] + timeoutSeconds: 600 + resources: + requests: + cpu: 100m + memory: 128Mi + actions: + - name: default + description: Run all diagnostics +`) + if err := yaml.Unmarshal(input, &decoded); err != nil { + t.Fatalf("decode manifest triage action: %v", err) + } + triage := decoded.Applications["weave-trace"].Triage + if triage.ContainerName != "weave-trace" || triage.TimeoutSeconds != 600 || + len(triage.Actions) != 1 || triage.Actions[0].Name != "default" { + t.Fatalf("decoded triage = %#v", triage) + } +} + +type staticActionLogReader struct { + output []byte + err error +} + +func (r *staticActionLogReader) ReadPodLogs( + _ context.Context, + _ string, + _ string, + _ string, + _ int64, +) ([]byte, error) { + return r.output, r.err +} + +func newActionTestScheme(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 testActionRun() *wandbv2.ActionRun { + return &wandbv2.ActionRun{ + TypeMeta: metav1.TypeMeta{ + APIVersion: wandbv2.GroupVersion.String(), + Kind: "ActionRun", + }, + ObjectMeta: metav1.ObjectMeta{ + Name: "weave-check", + Namespace: "wandb", + UID: types.UID("action-run-uid"), + }, + Spec: wandbv2.ActionRunSpec{ + Type: wandbv2.ActionTypeTriage, + ApplicationRef: wandbv2.ApplicationReference{Name: "weave-trace"}, + Action: wandbv2.ActionReference{Name: "default"}, + }, + } +} + +func testActionApplication() *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{ + ContainerName: "weave-trace", + Args: []string{"python", "-m", "weave_triage"}, + Env: []corev1.EnvVar{ + {Name: "PYTHONPATH", Value: "/weave/src"}, + {Name: actionTypeEnv, Value: "manifest-value-must-not-win"}, + {Name: actionNameEnv, Value: "manifest-value-must-not-win"}, + }, + Actions: []wandbv2.ApplicationActionSpec{{ + Name: "default", + Description: "Run all diagnostics", + }}, + }, + }, + } +} + +func requestFor(run *wandbv2.ActionRun) 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) +} + +func conditionReason(conditions []metav1.Condition, conditionType string) string { + for _, condition := range conditions { + if condition.Type == conditionType { + return condition.Reason + } + } + return "" +} diff --git a/internal/controller/reconciler/triage.go b/internal/controller/reconciler/triage.go index 98145e1b..b0c67936 100644 --- a/internal/controller/reconciler/triage.go +++ b/internal/controller/reconciler/triage.go @@ -15,10 +15,10 @@ func resolveApplicationTriage(triage *serverManifest.ApplicationTriage) *apiv2.A for i := range triage.Env { env[i] = *triage.Env[i].DeepCopy() } - actions := make([]apiv2.TriageActionSpec, len(triage.Actions)) + actions := make([]apiv2.ApplicationActionSpec, len(triage.Actions)) for i := range triage.Actions { - actions[i] = apiv2.TriageActionSpec{ - Name: apiv2.TriageActionName(triage.Actions[i].Name), + actions[i] = apiv2.ApplicationActionSpec{ + Name: apiv2.ActionName(triage.Actions[i].Name), Description: triage.Actions[i].Description, Args: append([]string(nil), triage.Actions[i].Args...), } diff --git a/internal/controller/triagerun_controller.go b/internal/controller/triagerun_controller.go deleted file mode 100644 index 539315f3..00000000 --- a/internal/controller/triagerun_controller.go +++ /dev/null @@ -1,889 +0,0 @@ -/* -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" - "strconv" - "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" - triageActionAnnotation = "apps.wandb.com/triage-action" - triageActionEnv = "WANDB_TRIAGE_ACTION" -) - -// 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 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 - 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 - } - - actions, err := resolveRequestedTriageActions(&run, &application) - if err != nil { - return r.failRun(ctx, &run, "InvalidAction", 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) - } - } - - 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 - } - } - if requeueForOutput { - return ctrl.Result{RequeueAfter: 2 * time.Second}, nil - } - return ctrl.Result{}, nil -} - -type resolvedTriageAction struct { - name wandbv2.TriageActionName - runner *wandbv2.ApplicationTriageSpec - 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, actionRef := range run.Spec.Actions { - actionName := actionRef.Name - spec, err := resolveTriageAction(application, actionName) - if err != nil { - return nil, err - } - runner := application.Spec.Triage - source, err := selectTriageContainer(application, runner.ContainerName) - if err != nil { - return nil, err - } - 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, runner, 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: action.jobName}, &job) - if err != nil && !apierrors.IsNotFound(err) { - return nil, err - } - if apierrors.IsNotFound(err) { - if previous != nil && previous.JobRef != nil { - return nil, &triageJobMissingError{ - name: previous.JobRef.Name, action: action.name, - } - } - 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 - } - if err := r.Create(ctx, &job); err != nil { - if !apierrors.IsAlreadyExists(err) { - return nil, err - } - if err := r.Get(ctx, types.NamespacedName{ - Namespace: run.Namespace, - Name: action.jobName, - }, &job); err != nil { - return nil, err - } - } - } - return &job, nil -} - -type triageJobMissingError struct { - name string - action wandbv2.TriageActionName -} - -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 wandbv2.TriageActionName, -) (wandbv2.TriageActionSpec, error) { - if application.Spec.Triage == nil { - return wandbv2.TriageActionSpec{}, fmt.Errorf("Application %q does not declare triage actions", application.Name) - } - if len(application.Spec.Triage.Command) == 0 && len(application.Spec.Triage.Args) == 0 { - return wandbv2.TriageActionSpec{}, fmt.Errorf( - "Application %q triage runner must override command or args", application.Name) - } - for i := range application.Spec.Triage.Actions { - action := application.Spec.Triage.Actions[i] - if action.Name == actionName { - 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) { - 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, - runner *wandbv2.ApplicationTriageSpec, - 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, runner, 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, - runner *wandbv2.ApplicationTriageSpec, - 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, runner.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(runner.Command) > 0 { - container.Command = append([]string(nil), runner.Command...) - container.Args = nil - } - if len(runner.Args) > 0 { - container.Args = append([]string(nil), runner.Args...) - } - container.Args = append(container.Args, action.Args...) - container.Env = mergeTriageEnv(container.Env, []corev1.EnvVar{{ - Name: triageActionEnv, - Value: string(action.Name), - }}) - if runner.Resources != nil { - container.Resources = *runner.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, - runner *wandbv2.ApplicationTriageSpec, - action wandbv2.TriageActionSpec, - timeoutSeconds int64, -) *wandbv2.TriageResolvedExecution { - container := buildTriageContainer(source, runner, 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) actionStatus( - ctx context.Context, - action *resolvedTriageAction, - job *batchv1.Job, - 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) - } - if !jobComplete(job) { - return status, false, "" - } - - 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) - } - status.Phase = wandbv2.TriageRunPhaseSucceeded - status.CompletedAt = triageCompletionTime(job) - status.Results = results - status.Summary = summarizeTriageResults(results) - return status, false, "" -} - -func setAggregateTriageStatus(run *wandbv2.TriageRun, failureMessages []string) { - run.Status.ObservedGeneration = run.Generation - 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 from %d actions", totalResults, completedActions), - }) -} - -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) failRun( - ctx context.Context, - run *wandbv2.TriageRun, - reason string, - message string, -) (ctrl.Result, error) { - run.Status.Phase = wandbv2.TriageRunPhaseFailed - run.Status.ObservedGeneration = run.Generation - now := metav1.Now() - run.Status.CompletedAt = &now - 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 triageActionStartTime( - previous *wandbv2.TriageActionStatus, - job *batchv1.Job, -) *metav1.Time { - if job.Status.StartTime != nil { - return job.Status.StartTime.DeepCopy() - } - if previous != nil && previous.StartedAt != nil { - return previous.StartedAt.DeepCopy() - } - if !job.CreationTimestamp.IsZero() { - return job.CreationTimestamp.DeepCopy() - } - now := metav1.Now() - return &now -} - -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 deleted file mode 100644 index bf4a6c31..00000000 --- a/internal/controller/triagerun_controller_unit_test.go +++ /dev/null @@ -1,475 +0,0 @@ -package controller - -import ( - "context" - "fmt" - "slices" - "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() - logReader := &staticTriageLogReader{} - reconciler := &TriageRunReconciler{ - Client: fakeClient, - Scheme: testScheme, - PodLogs: logReader, - } - - _, 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-0", - }, &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 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) - } - if got := envValue(container.Env, "DATABASE_URL"); got != "mysql://wandb" { - t.Fatalf("DATABASE_URL = %q, want inherited env", got) - } - if got := envValue(container.Env, triageActionEnv); got != defaultTriageAction { - t.Fatalf("%s = %q, want %q", triageActionEnv, got, defaultTriageAction) - } - 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 len(updatedRun.Status.ActionStatuses) != 1 { - t.Fatalf("action statuses = %d, want 1", len(updatedRun.Status.ActionStatuses)) - } - 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.TriageActionReference{ - {Name: "default"}, - {Name: "deep"}, - } - application := testTriageApplication() - application.Spec.Triage.Actions = append(application.Spec.Triage.Actions, - wandbv2.TriageActionSpec{ - Name: "deep", - Description: "Run deeper diagnostics", - Args: []string{"--verbose"}, - }) - 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) - } - args := job.Spec.Template.Spec.Containers[0].Args - want := []string{"python", "-m", "weave_triage"} - if action != defaultTriageAction { - want = append(want, "--verbose") - } - if !slices.Equal(args, want) { - t.Fatalf("Job %q args = %#v, want %#v", job.Name, args, want) - } - if got := envValue(job.Spec.Template.Spec.Containers[0].Env, triageActionEnv); got != action { - t.Fatalf("Job %q %s = %q, want %q", job.Name, triageActionEnv, got, 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.Name { - t.Fatalf("action status %d = %q, want %q", - i, updatedRun.Status.ActionStatuses[i].Action, action.Name) - } - } - - 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) - } - } -} - -func TestTriageRunCollectsFailedCheckAsSuccessfulExecution(t *testing.T) { - t.Parallel() - - testScheme := newTriageTestScheme(t) - run := testTriageRun() - run.Status.Phase = wandbv2.TriageRunPhaseRunning - run.Status.ActionStatuses = []wandbv2.TriageActionStatus{{ - Action: defaultTriageAction, - JobRef: &corev1.LocalObjectReference{Name: "weave-check-triage-0"}, - }} - application := testTriageApplication() - action := application.Spec.Triage.Actions[0] - source := &application.Spec.PodTemplate.Spec.Containers[0] - job := buildTriageJob( - 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) - } - 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.ActionStatuses) != 1 || - len(updatedRun.Status.ActionStatuses[0].Results) != 2 { - t.Fatalf("action statuses = %#v, want 2 results", updatedRun.Status.ActionStatuses) - } -} - -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: - containerName: weave-trace - args: [python, -m, weave_triage] - timeoutSeconds: 600 - resources: - requests: - cpu: 100m - memory: 128Mi - actions: - - - name: default - description: Run all diagnostics -`) - if err := yaml.Unmarshal(input, &decoded); err != nil { - t.Fatalf("decode manifest triage action: %v", err) - } - 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) - } -} - -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"}, - Actions: []wandbv2.TriageActionReference{{ - Name: 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{ - ContainerName: "weave-trace", - Args: []string{"python", "-m", "weave_triage"}, - Env: []corev1.EnvVar{ - {Name: "PYTHONPATH", Value: "/weave/src"}, - {Name: triageActionEnv, Value: "manifest-value-must-not-win"}, - }, - Actions: []wandbv2.TriageActionSpec{{ - Name: defaultTriageAction, - Description: "Run all diagnostics", - }}, - }, - }, - } -} - -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) -}