diff --git a/api/v2/application_types.go b/api/v2/application_types.go index c5cf1d9c..8ad7e5b3 100644 --- a/api/v2/application_types.go +++ b/api/v2/application_types.go @@ -67,11 +67,89 @@ 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 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. +type ApplicationTriageSpec 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. 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 + 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"` + + // 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. "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. + // +kubebuilder:validation:MaxLength=512 + // +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. + // +kubebuilder:validation:MaxItems=32 + // +optional + Args []string `json:"args,omitempty"` +} + // HTTPRouteTemplateSpec contains the fields needed to build a Gateway API HTTPRoute. type HTTPRouteTemplateSpec struct { ParentRefs []gatewayv1.ParentReference `json:"parentRefs"` diff --git a/api/v2/zz_generated.deepcopy.go b/api/v2/zz_generated.deepcopy.go index 93c5bd13..53a6af4e 100644 --- a/api/v2/zz_generated.deepcopy.go +++ b/api/v2/zz_generated.deepcopy.go @@ -150,6 +150,11 @@ func (in *ApplicationSpec) DeepCopyInto(out *ApplicationSpec) { (*in)[i].DeepCopyInto(&(*out)[i]) } } + if in.Triage != nil { + in, out := &in.Triage, &out.Triage + *out = new(ApplicationTriageSpec) + (*in).DeepCopyInto(*out) + } if in.HTTPRouteTemplate != nil { in, out := &in.HTTPRouteTemplate, &out.HTTPRouteTemplate *out = new(HTTPRouteTemplateSpec) @@ -231,6 +236,50 @@ func (in *ApplicationStatus) DeepCopy() *ApplicationStatus { return out } +// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. +func (in *ApplicationTriageSpec) DeepCopyInto(out *ApplicationTriageSpec) { + *out = *in + if in.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. func (in *CertManagerConfig) DeepCopyInto(out *CertManagerConfig) { *out = *in @@ -1494,6 +1543,26 @@ func (in *TriageActionReference) DeepCopy() *TriageActionReference { return out } +// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. +func (in *TriageActionSpec) DeepCopyInto(out *TriageActionSpec) { + *out = *in + if in.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 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..3450d383 100644 --- a/config/crd/bases/apps.wandb.com_applications.yaml +++ b/config/crd/bases/apps.wandb.com_applications.yaml @@ -12801,6 +12801,170 @@ spec: type: type: string type: object + triage: + properties: + actions: + items: + properties: + args: + items: + type: string + maxItems: 32 + type: array + description: + maxLength: 512 + type: string + 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: + 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 + 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..af7f7a61 100644 --- a/config/samples/apps_v2_application.yaml +++ b/config/samples/apps_v2_application.yaml @@ -16,6 +16,21 @@ spec: - name: my-app image: nginx:1.14.2 + triage: + containerName: my-app + command: ["/bin/echo"] + timeoutSeconds: 60 + actions: + - name: default + description: Run the default application diagnostic + resources: + requests: + cpu: 25m + memory: 32Mi + limits: + cpu: 100m + memory: 128Mi + jobs: - metadata: name: my-job @@ -42,4 +57,3 @@ spec: - name: my-cronjob image: nginx:1.14.2 command: ["sleep", "10"] - diff --git a/deploy/operator/templates/wandb-operator-wandb-role.yaml b/deploy/operator/templates/wandb-operator-wandb-role.yaml index 290ae4e8..75a26d83 100644 --- a/deploy/operator/templates/wandb-operator-wandb-role.yaml +++ b/deploy/operator/templates/wandb-operator-wandb-role.yaml @@ -17,6 +17,14 @@ rules: - patch - update - watch + - apiGroups: + - apps.wandb.com + resources: + - triageruns + verbs: + - get + - list + - watch - apiGroups: - apps.wandb.com resources: @@ -28,6 +36,7 @@ rules: - apps.wandb.com resources: - applications/status + - triageruns/status - weightsandbiases/status verbs: - get @@ -210,4 +219,4 @@ subjects: - kind: ServiceAccount name: {{ include "wandb-operator.fullname" . }} namespace: {{ .Release.Namespace }} -{{- end }} \ No newline at end of file +{{- end }} diff --git a/internal/controller/reconciler/reconcile_v2.go b/internal/controller/reconciler/reconcile_v2.go index 57891d89..f9092eed 100644 --- a/internal/controller/reconciler/reconcile_v2.go +++ b/internal/controller/reconciler/reconcile_v2.go @@ -668,6 +668,7 @@ func reconcileApplications( setCustomCACertsChecksumAnnotation(&application.Spec.PodTemplate, caChecksum) application.Spec.HpaTemplate = ResolveAutoscaling(app, wandb) + application.Spec.Triage = resolveApplicationTriage(app.Triage) // Set shared service account for all W&B applications application.Spec.PodTemplate.Spec.ServiceAccountName = serviceAccountName diff --git a/internal/controller/reconciler/triage.go b/internal/controller/reconciler/triage.go new file mode 100644 index 00000000..98145e1b --- /dev/null +++ b/internal/controller/reconciler/triage.go @@ -0,0 +1,38 @@ +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 + } + + 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...), + } + } + 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 new file mode 100644 index 00000000..423f46bd --- /dev/null +++ b/internal/controller/reconciler/triage_test.go @@ -0,0 +1,60 @@ +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{ + 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[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 resolved.Resources == nil || resolved.Resources.Requests.Memory().String() != "128Mi" { + t.Fatalf("resources = %#v", resolved.Resources) + } + if resolved.TimeoutSeconds != 600 { + t.Fatalf("timeoutSeconds = %d", resolved.TimeoutSeconds) + } + if action.Name != "default" || action.Description != "Run all diagnostics" || + len(action.Args) != 1 || action.Args[0] != "--verbose" { + t.Fatalf("action = %#v", action) + } + + input.Args[0] = "mutated" + input.Actions[0].Args[0] = "mutated" + resources.Requests[corev1.ResourceMemory] = resource.MustParse("1Gi") + 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 new file mode 100644 index 00000000..b1e09de4 --- /dev/null +++ b/internal/controller/triagerun_controller.go @@ -0,0 +1,888 @@ +/* +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" + triageActionFlag = "--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...) + } + 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() + } + 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 new file mode 100644 index 00000000..c3b68719 --- /dev/null +++ b/internal/controller/triagerun_controller_unit_test.go @@ -0,0 +1,469 @@ +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 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, triageActionFlag, action, "--verbose") + } + if !slices.Equal(args, want) { + t.Fatalf("Job %q args = %#v, want %#v", job.Name, args, want) + } + } + + 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", + }}, + 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) +} diff --git a/internal/crdinstaller/crds/operator/apps.wandb.com_applications.yaml b/internal/crdinstaller/crds/operator/apps.wandb.com_applications.yaml index 67be4bc9..3450d383 100644 --- a/internal/crdinstaller/crds/operator/apps.wandb.com_applications.yaml +++ b/internal/crdinstaller/crds/operator/apps.wandb.com_applications.yaml @@ -12801,6 +12801,170 @@ spec: type: type: string type: object + triage: + properties: + actions: + items: + properties: + args: + items: + type: string + maxItems: 32 + type: array + description: + maxLength: 512 + type: string + 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: + 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 + required: + - actions + type: object volumeClaimTemplates: items: properties: diff --git a/pkg/wandb/manifest/manifest.go b/pkg/wandb/manifest/manifest.go index 1cf06bd7..ad7649ab 100644 --- a/pkg/wandb/manifest/manifest.go +++ b/pkg/wandb/manifest/manifest.go @@ -175,6 +175,26 @@ 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 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 { + 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 {