diff --git a/api/v2/application_types.go b/api/v2/application_types.go index c5cf1d9c..46decaca 100644 --- a/api/v2/application_types.go +++ b/api/v2/application_types.go @@ -67,11 +67,86 @@ 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 ActionRun resources whose type is triage. + // +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. 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. + // +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. The + // selected action's Args are appended when starting the diagnostic runner. + // +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 []ApplicationActionSpec `json:"actions"` +} + +// 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 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 + // +optional + Description string `json:"description,omitempty"` + + // 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"` +} + // 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 a6c4d885..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) @@ -482,6 +507,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([]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(corev1.ResourceRequirements) + (*in).DeepCopyInto(*out) + } + if in.Actions != nil { + in, out := &in.Actions, &out.Actions + *out = make([]ApplicationActionSpec, 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 diff --git a/cmd/manager/main.go b/cmd/manager/main.go index 7edc029a..ca366624 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", "ActionRun") + os.Exit(1) + } + if err = (&controller.ActionRunReconciler{ + Client: mgr.GetClient(), + Scheme: mgr.GetScheme(), + PodLogs: &controller.KubernetesActionPodLogReader{ + CoreV1: kubernetesClient.CoreV1(), + }, + }).SetupWithManager(mgr); err != nil { + setupLog.Error(err, "unable to create controller", "controller", "ActionRun") + 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..93cf69fa 100644 --- a/config/rbac/role.yaml +++ b/config/rbac/role.yaml @@ -71,6 +71,24 @@ rules: - statefulsets/status verbs: - get +- apiGroups: + - apps.wandb.com + resources: + - actionruns + verbs: + - get + - list + - watch +- apiGroups: + - apps.wandb.com + resources: + - actionruns/status + - applications/status + - weightsandbiases/status + verbs: + - get + - patch + - update - apiGroups: - apps.wandb.com resources: @@ -91,15 +109,6 @@ rules: - weightsandbiases/finalizers verbs: - update -- apiGroups: - - apps.wandb.com - resources: - - applications/status - - weightsandbiases/status - verbs: - - get - - patch - - update - 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..49229ca3 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: + - actionruns + verbs: + - get + - list + - watch - apiGroups: - apps.wandb.com resources: @@ -28,6 +36,7 @@ rules: - apps.wandb.com resources: - applications/status + - actionruns/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/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/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..b0c67936 --- /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.ApplicationActionSpec, len(triage.Actions)) + for i := range triage.Actions { + actions[i] = apiv2.ApplicationActionSpec{ + Name: apiv2.ActionName(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/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 ce1cd96a..042ff8cb 100644 --- a/pkg/wandb/manifest/manifest.go +++ b/pkg/wandb/manifest/manifest.go @@ -179,6 +179,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 {