feat(weave): Add weave-trace (#157)

Co-authored-by: Justin Brooks <jsbroks@gmail.com>
wandb · Jul 19, 2024 · 0dfc408 · 0dfc408
1 parent 314ff0e
commit 0dfc408
Show file tree

Hide file tree

Showing 16 changed files with 528 additions and 25 deletions.
diff --git a/charts/operator-wandb/Chart.lock b/charts/operator-wandb/Chart.lock
@@ -8,6 +8,9 @@ dependencies:
 - name: weave
   repository: file://charts/weave
   version: 0.1.0
+- name: weave-trace
+  repository: file://charts/weave-trace
+  version: 0.1.0
 - name: parquet
   repository: file://charts/parquet
   version: 0.1.0
@@ -38,5 +41,5 @@ dependencies:
 - name: yace
   repository: file://charts/yace
   version: 0.1.0
-digest: sha256:c12c533a22b6f593a526e3060597ee1591bf5a2cac4ff03c588758e0dbc65d1e
-generated: "2024-06-24T16:15:00.442236-07:00"
+digest: sha256:bca2b6781737da6806e4485605cf9ce87b1428944b14cb88f082024cc3500bbd
+generated: "2024-07-18T01:17:04.532871-04:00"
diff --git a/charts/operator-wandb/Chart.yaml b/charts/operator-wandb/Chart.yaml
@@ -2,7 +2,7 @@ apiVersion: v2
 name: operator-wandb
 description: A Helm chart for deploying W&B to Kubernetes
 type: application
-version: 0.15.0
+version: 0.15.1
 appVersion: 1.0.0
 icon: https://wandb.ai/logo.svg
 
@@ -24,6 +24,10 @@ dependencies:
     version: "*.*.*"
     repository: file://charts/weave
     condition: weave.install
+  - name: weave-trace
+    version: "*.*.*"
+    repository: file://charts/weave-trace
+    condition: weave-trace.install
   - name: parquet
     version: "*.*.*"
     repository: file://charts/parquet

diff --git a/charts/operator-wandb/charts/app/templates/deployment.yaml b/charts/operator-wandb/charts/app/templates/deployment.yaml
@@ -114,6 +114,10 @@ spec:
               value: "http://{{ .Release.Name }}-parquet:8087"
             - name: PARQUET_ENABLED
               value: "true"
+            {{- if index .Values.global "weave-trace" "enabled" }}
+            - name: WEAVE_TRACES_ENABLED
+              value: "true"
+            {{- end }}
 
             {{- if ne (include "wandb.redis.password" .) "" }}
             - name: REDIS_PASSWORD

diff --git a/charts/operator-wandb/charts/weave-trace/.helmignore b/charts/operator-wandb/charts/weave-trace/.helmignore
@@ -0,0 +1,23 @@
+# Patterns to ignore when building packages.
+# This supports shell glob matching, relative path matching, and
+# negation (prefixed with !). Only one pattern per line.
+.DS_Store
+# Common VCS dirs
+.git/
+.gitignore
+.bzr/
+.bzrignore
+.hg/
+.hgignore
+.svn/
+# Common backup files
+*.swp
+*.bak
+*.tmp
+*.orig
+*~
+# Various IDEs
+.project
+.idea/
+*.tmproj
+.vscode/
diff --git a/charts/operator-wandb/charts/weave-trace/Chart.yaml b/charts/operator-wandb/charts/weave-trace/Chart.yaml
@@ -0,0 +1,5 @@
+apiVersion: v2
+name: weave-trace
+description: A Helm chart for Kubernetes
+version: 0.1.0
+appVersion: "1.0.0"
diff --git a/charts/operator-wandb/charts/weave-trace/templates/_helpers.tpl b/charts/operator-wandb/charts/weave-trace/templates/_helpers.tpl
@@ -0,0 +1,110 @@
+{{/* vim: set filetype=mustache: */}}
+
+{{/*
+Expand the name of the chart.
+*/}}
+{{- define "weaveTrace.name" -}}
+{{- default .Chart.Name .Values.nameOverride | trunc 63 | trimSuffix "-" }}
+{{- end }}
+
+{{/*
+Create a default fully qualified name for weave-trace. (Should be something like wandb-weave-trace)
+We truncate at 63 chars because some Kubernetes name fields are limited to this (by the DNS naming spec).
+If release name contains chart name it will be used as a full name.
+*/}}
+{{- define "weaveTrace.fullname" -}}
+{{- if .Values.fullnameOverride }}
+{{- .Values.fullnameOverride | trunc 63 | trimSuffix "-" }}
+{{- else }}
+{{- $name := default .Chart.Name .Values.nameOverride }}
+{{- if contains $name .Release.Name }}
+{{- .Release.Name | trunc 63 | trimSuffix "-" }}
+{{- else }}
+{{- printf "%s-%s" .Release.Name $name | trunc 63 | trimSuffix "-" }}
+{{- end }}
+{{- end }}
+{{- end }}
+
+{{/*
+Create a default fully qualified name for the weave-trace migration. (Should be something like wandb-weave-trace-migrate)
+We truncate at 63 chars because some Kubernetes name fields are limited to this (by the DNS naming spec).
+*/}}
+{{- define "weaveTraceMigrate.fullname" -}}
+{{ printf "%s-migrate" (include "weaveTrace.fullname" .) | trunc 63 | trimSuffix "-" }}
+{{- end }}
+
+
+{{/*
+Create chart name and version as used by the chart label.
+*/}}
+{{- define "weaveTrace.chart" -}}
+{{- printf "%s-%s" .Chart.Name .Chart.Version | replace "+" "_" | trunc 63 | trimSuffix "-" }}
+{{- end }}
+
+{{/*
+Common labels
+*/}}
+{{- define "weaveTrace.labels" -}}
+helm.sh/chart: {{ include "weaveTrace.chart" . }}
+{{ include "weaveTrace.selectorLabels" . }}
+{{- if .Chart.AppVersion }}
+app.kubernetes.io/version: {{ .Chart.AppVersion | quote }}
+{{- end }}
+wandb.com/app-name: {{ include "weaveTrace.chart" . }}
+app.kubernetes.io/managed-by: {{ .Release.Service }}
+{{- end }}
+
+{{/*
+Selector labels
+*/}}
+{{- define "weaveTrace.selectorLabels" -}}
+app.kubernetes.io/name: {{ include "weaveTrace.name" . }}
+app.kubernetes.io/instance: {{ .Release.Name }}
+{{- end }}
+
+{{/*
+Create the name of the service account to use
+*/}}
+{{- define "weaveTrace.serviceAccountName" -}}
+{{- if .Values.serviceAccount.create }}
+{{- default (include "weaveTrace.fullname" .) .Values.serviceAccount.name }}
+{{- else }}
+{{- default "default" .Values.serviceAccount.name }}
+{{- end }}
+{{- end }}
+
+{{/*
+Returns the extraEnv keys and values to inject into containers.
+
+Global values will override any chart-specific values.
+*/}}
+{{- define "weaveTrace.extraEnv" -}}
+{{- $allExtraEnv := merge (default (dict) .local.extraEnv) .global.extraEnv -}}
+{{- range $key, $value := $allExtraEnv }}
+- name: {{ $key }}
+  value: {{ $value | quote }}
+{{- end -}}
+{{- end -}}
+
+{{/*
+Returns a list of _common_ labels to be shared across all
+app deployments and other shared objects.
+*/}}
+{{- define "weaveTrace.commonLabels" -}}
+{{- $commonLabels := default (dict) .Values.common.labels -}}
+{{- if $commonLabels }}
+{{-   range $key, $value := $commonLabels }}
+{{ $key }}: {{ $value | quote }}
+{{-   end }}
+{{- end -}}
+{{- end -}}
+
+{{/*
+Returns a list of _pod_ labels to be shared across all
+app deployments.
+*/}}
+{{- define "weaveTrace.podLabels" -}}
+{{- range $key, $value := .Values.pod.labels }}
+{{ $key }}: {{ $value | quote }}
+{{- end }}
+{{- end -}}
diff --git a/charts/operator-wandb/charts/weave-trace/templates/deployment.yaml b/charts/operator-wandb/charts/weave-trace/templates/deployment.yaml
@@ -0,0 +1,168 @@
+{{- $imageCfg := dict "global" $.Values.global.image "local" $.Values.image -}}
+apiVersion: apps/v1
+kind: Deployment
+metadata:
+  name: {{ include "weaveTrace.fullname" . }}
+  labels:
+    {{- include "wandb.commonLabels" . | nindent 4 }}
+    {{- include "weaveTrace.commonLabels" . | nindent 4 }}
+    {{- include "weaveTrace.labels" . | nindent 4 }}
+    {{- if .Values.deployment.labels -}}
+    {{-   toYaml .Values.deployment.labels | nindent 4 }}
+    {{- end }}
+  annotations:
+    {{- include "wandb.deploymentAnnotations" $ | nindent 4 }}
+    {{- if .Values.deployment.annotations -}}
+    {{-   toYaml .Values.deployment.annotations | nindent 4 }}
+    {{- end }}
+spec:
+  replicas: 1
+  selector:
+    matchLabels:
+      {{- include "wandb.selectorLabels" $ | nindent 6 }}
+      {{- include "weaveTrace.labels" . | nindent 6 }}
+  template:
+    metadata:
+      labels:
+        {{- include "wandb.podLabels" . | nindent 8 }}
+        {{- include "weaveTrace.commonLabels" . | nindent 8 }}
+        {{- include "weaveTrace.podLabels" . | nindent 8 }}
+        {{- include "weaveTrace.labels" . | nindent 8 }}
+      annotations:
+        {{- if .Values.pod.annotations -}}
+        {{-   toYaml .Values.pod.annotations | nindent 4 }}
+        {{- end }}
+    spec:
+      {{- if .tolerations }}
+      tolerations:
+        {{- toYaml .tolerations | nindent 8 }}
+      {{- end }}
+      {{- include "wandb.nodeSelector" . | nindent 6 }}
+      {{- include "wandb.priorityClassName" . | nindent 6 }}
+      {{- include "wandb.podSecurityContext" .Values.pod.securityContext | nindent 6 }}
+      {{- if not .Values.migration.useHook }}
+      initContainers:
+        - name: {{ include "weaveTraceMigrate.fullname" . }}
+          image: "{{ .Values.image.repository }}:{{ .Values.image.tag }}"
+          command:
+            - "python"
+            - "migrator.py"
+          env:
+            - name: WF_CLICKHOUSE_HOST
+              value: "{{ .Values.global.clickhouse.host }}"
+            - name: WF_CLICKHOUSE_PORT
+              value: "{{ .Values.global.clickhouse.port }}"
+            - name: WF_CLICKHOUSE_DATABASE
+              value: "{{ .Values.global.clickhouse.database }}"
+            - name: WF_CLICKHOUSE_USER
+              value: "{{ .Values.global.clickhouse.user }}"
+            - name: WF_CLICKHOUSE_PASS
+              valueFrom:
+                secretKeyRef:
+                  name: {{ include "wandb.clickhouse.passwordSecret" . }}
+                  key: CLICKHOUSE_PASSWORD
+            {{- include "weaveTrace.extraEnv" (dict "global" .Values.global "local" .Values) | nindent 12 }}
+            {{- include "wandb.extraEnvFrom" (dict "root" $ "local" .) | nindent 12 }}
+      {{- end }}
+      containers:
+        - name: {{ include "weaveTrace.fullname" .  }}
+          image: "{{ .Values.image.repository }}:{{ .Values.image.tag }}"
+          ports:
+            - name: http
+              containerPort: 8080
+              protocol: TCP
+          env:
+            - name: PORT
+              value: "8080"
+            - name: API_PATH_PREFIX
+              value: "/traces"
+            - name: WANDB_PUBLIC_BASE_URL
+              value: {{ .Values.global.host }}
+            - name: WANDB_BASE_URL
+              value: http://{{ .Release.Name }}-app:8080/
+            - name: WF_TRACE_SERVER_URL
+              value: "{{ .Values.global.host }}/traces"
+            - name: WF_ENFORCE_PASSWORD_LENGTH
+              value: "false"
+            - name: WF_CLICKHOUSE_HOST
+              value: "{{ .Values.global.clickhouse.host }}"
+            - name: WF_CLICKHOUSE_PORT
+              value: "{{ .Values.global.clickhouse.port }}"
+            - name: WF_CLICKHOUSE_DATABASE
+              value: "{{ .Values.global.clickhouse.database }}"
+            - name: WF_CLICKHOUSE_USER
+              value: "{{ .Values.global.clickhouse.user }}"
+            - name: WF_CLICKHOUSE_PASS
+              valueFrom:
+                secretKeyRef:
+                  name: {{ include "wandb.clickhouse.passwordSecret" . }}
+                  key: CLICKHOUSE_PASSWORD
+            {{- if .Values.datadog.enabled }}
+            - name: DD_SERVICE
+              value: "{{ .Values.datadog.service }}"
+            - name: DD_ENV
+              value: "{{ .Values.datadog.env }}"
+            - name: DD_TRACE_ENABLED
+              value: "{{ .Values.datadog.traceEnabled }}"
+            - name: DD_LOGS_ENABLED
+              value: "{{ .Values.env.logsEnabled }}"
+            - name: DD_LOGS_INJECTION
+              value: "{{ .Values.env.logsInjection }}"
+            {{- end }}
+            {{- include "weaveTrace.extraEnv" (dict "global" .Values.global "local" .Values) | nindent 12 }}
+            {{- include "wandb.extraEnvFrom" (dict "root" $ "local" .) | nindent 12 }}
+          {{- if not .Values.datadog.enabled }}
+          command:
+            - uvicorn
+          args:
+            - "src.trace_server:app"
+            - "--host"
+            - "0.0.0.0"
+            - "--port"
+            - "8080"
+          {{- end }}
+          livenessProbe:
+            httpGet:
+              path: /traces/health
+              port: http
+            timeoutSeconds: 2
+            failureThreshold: 5
+          readinessProbe:
+            httpGet:
+              path: /traces/health
+              port: http
+            timeoutSeconds: 2
+            failureThreshold: 5
+          startupProbe:
+            httpGet:
+              path: /traces/health
+              port: http
+            failureThreshold: 12
+            periodSeconds: 10
+
+          resources:
+            {{- toYaml .Values.resources | nindent 12 }}
+---
+apiVersion: autoscaling/v2
+kind: HorizontalPodAutoscaler
+metadata:
+  name:  {{ include "weaveTrace.fullname" . }}
+  labels:
+    {{- include "wandb.commonLabels" . | nindent 4 }}
+    {{- include "weaveTrace.commonLabels" . | nindent 4 }}
+    {{- include "weaveTrace.labels" . | nindent 4 }}
+spec:
+  scaleTargetRef:
+    apiVersion: apps/v1
+    kind: Deployment
+    name: {{ include "weaveTrace.fullname" . }}
+  minReplicas: 1
+  maxReplicas: 5
+  metrics:
+  - type: Resource
+    resource:
+      name: cpu
+      target:
+        type: Utilization
+        averageUtilization: 70
+