Feature: pod annotation metadata

k8s/vizier/etcd_metadata/base/metadata_deployment.yaml

@@ -100,6 +100,8 @@ spec:
           value: "https://pl-etcd-client.$(PL_POD_NAMESPACE).svc:2379"
         - name: PL_ETCD_OPERATOR_ENABLED
           value: "true"
+        - name: PL_POD_ANNOTATION_ALLOWLIST
+          value: ""
         envFrom:
         - configMapRef:
             name: pl-tls-config

k8s/vizier/persistent_metadata/base/metadata_statefulset.yaml

@@ -84,6 +84,8 @@ spec:
           value: "7500"
         - name: PL_ETCD_OPERATOR_ENABLED
           value: "false"
+        - name: PL_POD_ANNOTATION_ALLOWLIST
+          value: ""
         envFrom:
         - configMapRef:
             name: pl-tls-config

src/carnot/funcs/metadata/metadata_ops.cc

@@ -49,6 +49,7 @@ void RegisterMetadataOpsOrDie(px::carnot::udf::Registry* registry) {
   registry->RegisterOrDie<IPToPodIDAtTimeUDF>("ip_to_pod_id");
   registry->RegisterOrDie<PodIDToPodNameUDF>("pod_id_to_pod_name");
   registry->RegisterOrDie<PodIDToPodLabelsUDF>("pod_id_to_pod_labels");
+  registry->RegisterOrDie<PodIDToPodAnnotationsUDF>("pod_id_to_pod_annotations");
   registry->RegisterOrDie<PodIDToNamespaceUDF>("pod_id_to_namespace");
   registry->RegisterOrDie<PodIDToNodeNameUDF>("pod_id_to_node_name");
   registry->RegisterOrDie<PodIDToReplicaSetNameUDF>("pod_id_to_replicaset_name");

src/carnot/funcs/metadata/metadata_ops.h

@@ -144,6 +144,27 @@ class PodIDToPodLabelsUDF : public ScalarUDF {
   }
 };
 
+class PodIDToPodAnnotationsUDF : public ScalarUDF {
+ public:
+  StringValue Exec(FunctionContext* ctx, StringValue pod_id) {
+    auto md = GetMetadataState(ctx);
+
+    const auto* pod_info = md->k8s_metadata_state().PodInfoByID(pod_id);
+    if (pod_info != nullptr) {
+      return pod_info->annotations();
+    }
+    return "";
+  }
+
+  static udf::ScalarUDFDocBuilder Doc() {
+    return udf::ScalarUDFDocBuilder("Get annotations of a pod from its pod ID.")
+        .Details("Gets the kubernetes pod annotations for the pod from its pod ID.")
+        .Example("df.annotations = px.pod_id_to_pod_annotations(df.pod_id)")
+        .Arg("pod_id", "The pod ID of the pod to get the annotations for.")
+        .Returns("The k8s pod annotations for the pod ID passed in.");
+  }
+};
+
 class PodNameToPodIDUDF : public ScalarUDF {
  public:
   StringValue Exec(FunctionContext* ctx, StringValue pod_name) {

src/carnot/funcs/metadata/metadata_ops_test.cc

@@ -124,6 +124,14 @@ TEST_F(MetadataOpsTest, pod_id_to_pod_labels_test) {
   udf_tester.ForInput("missing").Expect("");
 }
 
+TEST_F(MetadataOpsTest, pod_id_to_pod_annotations_test) {
+  auto function_ctx = std::make_unique<FunctionContext>(metadata_state_, nullptr);
+  auto udf_tester = px::carnot::udf::UDFTester<PodIDToPodAnnotationsUDF>(std::move(function_ctx));
+  udf_tester.ForInput("1_uid").Expect("{\"a1\":\"v1\", \"a2\":\"v2\"}");
+  udf_tester.ForInput("2_uid").Expect("{\"a1\":\"v1\"}");
+  udf_tester.ForInput("missing").Expect("");
+}
+
 TEST_F(MetadataOpsTest, pod_name_to_pod_id_test) {
   auto function_ctx = std::make_unique<FunctionContext>(metadata_state_, nullptr);
   auto udf_tester = px::carnot::udf::UDFTester<PodNameToPodIDUDF>(std::move(function_ctx));

src/shared/k8s/metadatapb/metadata.proto

@@ -414,6 +414,8 @@ message PodUpdate {
   // A json object containing pod label keys and values
   string labels = 17;
   repeated OwnerReference owner_references = 18;
+  // A json object containing pod annotation keys and values
+  string annotations = 19;
 }
 
 enum ContainerType {

src/shared/k8s/metadatapb/test_proto.h

@@ -50,6 +50,7 @@ uid: "1_uid"
 name: "running_pod"
 namespace: "pl"
 labels: "{\"k1\":\"v1\", \"k2\":\"v2\"}"
+annotations: "{\"a1\":\"v1\", \"a2\":\"v2\"}"
 start_timestamp_ns: 5
 container_ids: "pod1_container_1"
 qos_class: QOS_CLASS_GUARANTEED
@@ -94,6 +95,7 @@ uid: "2_uid"
 name: "terminating_pod"
 namespace: "pl"
 labels: "{\"k1\":\"v1\"}"
+annotations: "{\"a1\":\"v1\"}"
 start_timestamp_ns: 10
 container_ids: "pod2_container_1"
 qos_class: QOS_CLASS_BEST_EFFORT

src/shared/metadata/k8s_objects.h

@@ -323,11 +323,13 @@ class PodInfo : public K8sMetadataObject {
   void set_hostname(std::string_view hostname) { hostname_ = hostname; }
   void set_pod_ip(std::string_view pod_ip) { pod_ip_ = pod_ip; }
   void set_pod_labels(std::string labels) { labels_ = labels; }
+  void set_pod_annotations(std::string annotations) { annotations_ = annotations; }
 
   const std::string& node_name() const { return node_name_; }
   const std::string& hostname() const { return hostname_; }
   const std::string& pod_ip() const { return pod_ip_; }
   const std::string& labels() const { return labels_; }
+  const std::string& annotations() const { return annotations_; }
 
   const absl::flat_hash_set<std::string>& containers() const { return containers_; }
   const absl::flat_hash_set<std::string>& services() const { return services_; }
@@ -368,6 +370,7 @@ class PodInfo : public K8sMetadataObject {
   std::string hostname_;
   std::string pod_ip_;
   std::string labels_;
+  std::string annotations_;
 };
 
 struct UPIDStartTSCompare {

src/shared/metadata/metadata_state.cc

@@ -292,6 +292,7 @@ Status K8sMetadataState::HandlePodUpdate(const PodUpdate& update) {
   pod_info->set_phase_message(update.message());
   pod_info->set_phase_reason(update.reason());
   pod_info->set_pod_labels(update.labels());
+  pod_info->set_pod_annotations(update.annotations());
 
   pods_by_name_[{ns, name}] = object_uid;
   // Filter out daemonsets which don't have their own, unique podIP.

src/shared/metadata/metadata_state_test.cc

@@ -34,6 +34,7 @@ constexpr char kPod0UpdatePbTxt[] = R"(
   name: "pod0"
   namespace: "ns0"
   labels: "{\"k1\":\"v1\", \"k2\":\"v2\"}"
+  annotations: "{\"a1\":\"v1\"}"
   start_timestamp_ns: 101
   stop_timestamp_ns: 103
   container_ids: "container0_uid"
@@ -373,6 +374,7 @@ TEST(K8sMetadataStateTest, HandlePodUpdate) {
   EXPECT_EQ("pod0", pod_info->name());
   EXPECT_EQ("ns0", pod_info->ns());
   EXPECT_EQ("{\"k1\":\"v1\", \"k2\":\"v2\"}", pod_info->labels());
+  EXPECT_EQ("{\"a1\":\"v1\"}", pod_info->annotations());
   EXPECT_EQ(PodQOSClass::kGuaranteed, pod_info->qos_class());
   EXPECT_EQ(PodPhase::kRunning, pod_info->phase());
   EXPECT_EQ(1, pod_info->conditions().size());

src/vizier/services/metadata/controllers/k8smeta/k8s_metadata_handler.go

@@ -173,17 +173,24 @@ type Handler struct {
 }
 
 // NewHandler creates a new Handler.
-func NewHandler(updateCh <-chan *K8sResourceMessage, mds Store, pls PodLabelStore, conn *nats.Conn) *Handler {
+func NewHandler(updateCh <-chan *K8sResourceMessage, mds Store, pls PodLabelStore, conn *nats.Conn, podAnnotationAllowlist []string) *Handler {
 	done := make(chan struct{})
 	leaderMsgs := make(map[string]*metadatapb.Endpoints)
 	handlerMap := make(map[string]UpdateProcessor)
 	state := ProcessorState{LeaderMsgs: leaderMsgs, PodCIDRs: make([]string, 0), NodeToIP: make(map[string]string), PodToIP: make(map[string]string)}
 	mh := &Handler{updateCh: updateCh, mds: mds, pls: pls, conn: conn, done: done, processHandlerMap: handlerMap, state: state}
 
+	annotationAllowlist := make(map[string]bool, len(podAnnotationAllowlist))
+	for _, k := range podAnnotationAllowlist {
+		if k != "" {
+			annotationAllowlist[k] = true
+		}
+	}
+
 	// Register update processors.
 	mh.processHandlerMap["endpoints"] = &EndpointsUpdateProcessor{}
 	mh.processHandlerMap["services"] = &ServiceUpdateProcessor{}
-	mh.processHandlerMap["pods"] = &PodUpdateProcessor{}
+	mh.processHandlerMap["pods"] = &PodUpdateProcessor{annotationAllowlist: annotationAllowlist}
 	mh.processHandlerMap["nodes"] = &NodeUpdateProcessor{}
 	mh.processHandlerMap["namespaces"] = &NamespaceUpdateProcessor{}
 	mh.processHandlerMap["replicasets"] = &ReplicaSetUpdateProcessor{}
@@ -627,7 +634,10 @@ func (p *ServiceUpdateProcessor) GetUpdatesToSend(storedUpdates []*StoredUpdate,
 }
 
 // PodUpdateProcessor is a processor for pods.
-type PodUpdateProcessor struct{}
+type PodUpdateProcessor struct {
+	// The set of pod annotation keys to capture. Annotations not in this set are dropped.
+	annotationAllowlist map[string]bool
+}
 
 // IsNodeScoped returns whether this update is scoped to specific nodes, or should be sent to all nodes.
 func (p *PodUpdateProcessor) IsNodeScoped() bool {
@@ -737,7 +747,7 @@ func (p *PodUpdateProcessor) GetUpdatesToSend(storedUpdates []*StoredUpdate, sta
 		podUpdate := u.Update.GetPod()
 		if podUpdate != nil {
 			updates = append(updates, &OutgoingUpdate{
-				Update: getResourceUpdateFromPod(podUpdate, u.UpdateVersion),
+				Update: getResourceUpdateFromPod(podUpdate, u.UpdateVersion, p.annotationAllowlist),
 				Topics: topics,
 			})
 		}
@@ -1087,7 +1097,7 @@ func getServiceResourceUpdateFromEndpoint(ep *metadatapb.Endpoints, uv int64, po
 	return update
 }
 
-func getResourceUpdateFromPod(pod *metadatapb.Pod, uv int64) *metadatapb.ResourceUpdate {
+func getResourceUpdateFromPod(pod *metadatapb.Pod, uv int64, annotationAllowlist map[string]bool) *metadatapb.ResourceUpdate {
 	var containerIDs []string
 	var containerNames []string
 	if pod.Status.ContainerStatuses != nil {
@@ -1109,6 +1119,19 @@ func getResourceUpdateFromPod(pod *metadatapb.Pod, uv int64) *metadatapb.Resourc
 		podLabels, _ = json.Marshal(pod.Metadata.Labels)
 	}
 
+	var podAnnotations []byte
+	if len(annotationAllowlist) > 0 && pod.Metadata.Annotations != nil {
+		filtered := make(map[string]string)
+		for k, v := range pod.Metadata.Annotations {
+			if annotationAllowlist[k] {
+				filtered[k] = v
+			}
+		}
+		if len(filtered) > 0 {
+			podAnnotations, _ = json.Marshal(filtered)
+		}
+	}
+
 	update := &metadatapb.ResourceUpdate{
 		UpdateVersion: uv,
 		Update: &metadatapb.ResourceUpdate_PodUpdate{
@@ -1117,6 +1140,7 @@ func getResourceUpdateFromPod(pod *metadatapb.Pod, uv int64) *metadatapb.Resourc
 				Name:             pod.Metadata.Name,
 				Namespace:        pod.Metadata.Namespace,
 				Labels:           string(podLabels),
+				Annotations:      string(podAnnotations),
 				StartTimestampNS: pod.Metadata.CreationTimestampNS,
 				StopTimestampNS:  pod.Metadata.DeletionTimestampNS,
 				QOSClass:         pod.Status.QOSClass,

src/vizier/services/metadata/controllers/k8smeta/k8s_metadata_handler_test.go

@@ -23,6 +23,7 @@ import (
 	"sort"
 	"sync"
 	"testing"
+	"time"
 
 	"github.com/gogo/protobuf/proto"
 	"github.com/nats-io/nats.go"
@@ -399,7 +400,7 @@ func TestHandler_GetUpdatesForIP(t *testing.T) {
 	require.NoError(t, err)
 
 	updateCh := make(chan *k8smeta.K8sResourceMessage)
-	mdh := k8smeta.NewHandler(updateCh, mds, lps, nil)
+	mdh := k8smeta.NewHandler(updateCh, mds, lps, nil, nil)
 	defer mdh.Stop()
 	updates, err := mdh.GetUpdatesForIP("", 0, 0)
 	require.NoError(t, err)
@@ -448,7 +449,7 @@ func TestHandler_ProcessUpdates(t *testing.T) {
 	nc, natsCleanup := testingutils.MustStartTestNATS(t)
 	defer natsCleanup()
 
-	mdh := k8smeta.NewHandler(updateCh, mds, lps, nc)
+	mdh := k8smeta.NewHandler(updateCh, mds, lps, nc, nil)
 	defer mdh.Stop()
 
 	expectedNSMsg := &messagespb.VizierMessage{
@@ -624,6 +625,98 @@ func TestHandler_ProcessUpdates(t *testing.T) {
 	}, mds.ResourceStoreByTopic["unscoped"][5])
 }
 
+func TestHandler_PodAnnotationAllowlist(t *testing.T) {
+	tests := []struct {
+		name      string
+		allowlist []string
+		expected  string
+	}{
+		{
+			name:      "empty allowlist captures nothing",
+			allowlist: []string{},
+			expected:  "",
+		},
+		{
+			name:      "nil allowlist captures nothing",
+			allowlist: nil,
+			expected:  "",
+		},
+		{
+			name:      "only allowed keys are captured",
+			allowlist: []string{"team"},
+			expected:  `{"team":"pixie"}`,
+		},
+		{
+			name:      "non-matching key captures nothing",
+			allowlist: []string{"nonexistent"},
+			expected:  "",
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			updateCh := make(chan *k8smeta.K8sResourceMessage)
+			mds := &InMemoryStore{
+				ResourceStoreByTopic: make(map[string]ResourceStore),
+				RVStore:              map[string]int64{},
+				FullResourceStore:    make(map[int64]*storepb.K8SResource),
+			}
+			lps := &testutils.InMemoryPodLabelStore{
+				Store: make(map[string]string),
+			}
+
+			nc, natsCleanup := testingutils.MustStartTestNATS(t)
+			defer natsCleanup()
+
+			mdh := k8smeta.NewHandler(updateCh, mds, lps, nc, tc.allowlist)
+			defer mdh.Stop()
+
+			// Pod updates are always published to the Kelvin topic, so capture the
+			// emitted pod annotations from there.
+			gotCh := make(chan string, 1)
+			_, err := nc.Subscribe(fmt.Sprintf("%s/%s", k8smeta.K8sMetadataUpdateChannel, k8smeta.KelvinUpdateTopic), func(msg *nats.Msg) {
+				m := &messagespb.VizierMessage{}
+				if err := proto.Unmarshal(msg.Data, m); err != nil {
+					return
+				}
+				pu := m.GetK8SMetadataMessage().GetK8SMetadataUpdate().GetPodUpdate()
+				if pu != nil {
+					gotCh <- pu.Annotations
+				}
+			})
+			require.NoError(t, err)
+
+			updateCh <- &k8smeta.K8sResourceMessage{
+				ObjectType: "pods",
+				Object: &storepb.K8SResource{
+					Resource: &storepb.K8SResource_Pod{
+						Pod: &metadatapb.Pod{
+							Metadata: &metadatapb.ObjectMetadata{
+								UID:       "ijkl",
+								Name:      "object_md",
+								Namespace: "ns",
+								Annotations: map[string]string{
+									"team":                         "pixie",
+									"kubectl.kubernetes.io/loaded": "big-blob",
+								},
+							},
+							Status: &metadatapb.PodStatus{HostIP: "127.0.0.1"},
+							Spec:   &metadatapb.PodSpec{},
+						},
+					},
+				},
+			}
+
+			select {
+			case got := <-gotCh:
+				assert.Equal(t, tc.expected, got)
+			case <-time.After(5 * time.Second):
+				t.Fatal("timed out waiting for pod update")
+			}
+		})
+	}
+}
+
 func TestEndpointsUpdateProcessor_SetDeleted(t *testing.T) {
 	// Construct endpoints object.
 	o := createEndpointsObject()

src/vizier/services/metadata/controllers/k8smeta/metadata_topic_listener_test.go

@@ -136,7 +136,7 @@ func TestMetadataTopicListener_GetUpdatesInBatches(t *testing.T) {
 		t.Run(test.name, func(t *testing.T) {
 			mds := &FakeStore{}
 			updateCh := make(chan *K8sResourceMessage)
-			mdh := NewHandler(updateCh, mds, mds, nil)
+			mdh := NewHandler(updateCh, mds, mds, nil, nil)
 			mdTL, err := NewMetadataTopicListener(mdh, func(topic string, b []byte) error {
 				return nil
 			})
@@ -176,7 +176,7 @@ func TestMetadataTopicListener_GetUpdatesInBatches(t *testing.T) {
 func TestMetadataTopicListener_ProcessAgentMessage(t *testing.T) {
 	mds := &FakeStore{}
 	updateCh := make(chan *K8sResourceMessage)
-	mdh := NewHandler(updateCh, mds, mds, nil)
+	mdh := NewHandler(updateCh, mds, mds, nil, nil)
 
 	sentUpdates := make([]*messagespb.VizierMessage, 0)
 	mdTL, err := NewMetadataTopicListener(mdh, func(topic string, b []byte) error {

src/vizier/services/metadata/metadata_server.go

@@ -76,6 +76,7 @@ func init() {
 	pflag.String("nats_url", "pl-nats", "The URL of NATS")
 	pflag.Bool("use_etcd_operator", false, "Whether the etcd operator should be used instead of the persistent version.")
 	pflag.StringSlice("metadata_namespaces", []string{v1.NamespaceAll}, "The list of namespaces to watch for metadata.")
+	pflag.StringSlice("pod_annotation_allowlist", []string{}, "The list of pod annotation keys to capture. Empty captures none.")
 
 	// Metadata flags are set using the env vars in pl-cluster-config.
 	// We historically set PL_ETCD_OPERATOR_ENABLED but not PL_USE_ETCD_OPERATOR in the configmap.
@@ -236,7 +237,8 @@ func main() {
 	k8sMds := k8smeta.NewDatastore(dataStore)
 	// Listen for K8s metadata updates.
 	updateCh := make(chan *k8smeta.K8sResourceMessage)
-	mdh := k8smeta.NewHandler(updateCh, k8sMds, k8sMds, nc)
+	podAnnotationAllowlist := viper.GetStringSlice("pod_annotation_allowlist")
+	mdh := k8smeta.NewHandler(updateCh, k8sMds, k8sMds, nc, podAnnotationAllowlist)
 
 	namespaces := viper.GetStringSlice("metadata_namespaces")
 	if len(namespaces) == 0 {