Add PxL UDFs to look up the host IP of a pod

src/carnot/funcs/metadata/metadata_ops.cc

@@ -51,6 +51,7 @@ void RegisterMetadataOpsOrDie(px::carnot::udf::Registry* registry) {
   registry->RegisterOrDie<PodIDToPodLabelsUDF>("pod_id_to_pod_labels");
   registry->RegisterOrDie<PodIDToNamespaceUDF>("pod_id_to_namespace");
   registry->RegisterOrDie<PodIDToNodeNameUDF>("pod_id_to_node_name");
+  registry->RegisterOrDie<PodIDToHostIPUDF>("pod_id_to_host_ip");
   registry->RegisterOrDie<PodIDToReplicaSetNameUDF>("pod_id_to_replicaset_name");
   registry->RegisterOrDie<PodIDToReplicaSetIDUDF>("pod_id_to_replicaset_id");
   registry->RegisterOrDie<PodIDToDeploymentNameUDF>("pod_id_to_deployment_name");
@@ -68,6 +69,7 @@ void RegisterMetadataOpsOrDie(px::carnot::udf::Registry* registry) {
   registry->RegisterOrDie<PodNameToDeploymentIDUDF>("pod_name_to_deployment_id");
   registry->RegisterOrDie<PodNameToPodIDUDF>("pod_name_to_pod_id");
   registry->RegisterOrDie<PodNameToPodIPUDF>("pod_name_to_pod_ip");
+  registry->RegisterOrDie<PodNameToHostIPUDF>("pod_name_to_host_ip");
   registry->RegisterOrDie<PodNameToServiceNameUDF>("pod_name_to_service_name");
   registry->RegisterOrDie<PodNameToServiceIDUDF>("pod_name_to_service_id");
   registry->RegisterOrDie<PodNameToPodStartTimeUDF>("pod_name_to_start_time");

src/carnot/funcs/metadata/metadata_ops.h

@@ -188,6 +188,29 @@ class PodNameToPodIPUDF : public ScalarUDF {
   }
 };
 
+class PodNameToHostIPUDF : public ScalarUDF {
+ public:
+  StringValue Exec(FunctionContext* ctx, StringValue pod_name) {
+    auto md = GetMetadataState(ctx);
+    StringValue pod_id = PodNameToPodIDUDF::GetPodID(md, pod_name);
+    const px::md::PodInfo* pod_info = md->k8s_metadata_state().PodInfoByID(pod_id);
+    if (pod_info == nullptr) {
+      return "";
+    }
+    return pod_info->host_ip();
+  }
+
+  static udf::ScalarUDFDocBuilder Doc() {
+    return udf::ScalarUDFDocBuilder("Get the IP address of the node a pod is running on.")
+        .Details(
+            "Gets the IP address of the host (node) that the pod (specified by name) is "
+            "assigned to.")
+        .Example("df.host_ip = px.pod_name_to_host_ip(df.pod_name)")
+        .Arg("pod_name", "The name of the pod to get the host IP for.")
+        .Returns("The host IP for the pod name passed in.");
+  }
+};
+
 class PodIDToNamespaceUDF : public ScalarUDF {
  public:
   StringValue Exec(FunctionContext* ctx, StringValue pod_id) {
@@ -1945,6 +1968,29 @@ class PodIDToNodeNameUDF : public ScalarUDF {
   }
 };
 
+class PodIDToHostIPUDF : 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 "";
+    }
+    return pod_info->host_ip();
+  }
+
+  static udf::ScalarUDFDocBuilder Doc() {
+    return udf::ScalarUDFDocBuilder("Get the IP address of the node a pod ID is running on.")
+        .Details(
+            "Gets the IP address of the host (node) that the Pod (specified by Pod ID) is "
+            "assigned to.")
+        .Example("df.host_ip = px.pod_id_to_host_ip(df.pod_id)")
+        .Arg("pod_id", "The Pod ID of the Pod to get the host IP for.")
+        .Returns("The host IP for the Pod ID passed in.");
+  }
+};
+
 /**
  * @brief Returns the ReplicaSet name of a pod ID passed in.
  */

src/carnot/funcs/metadata/metadata_ops_test.cc

@@ -311,6 +311,14 @@ TEST_F(MetadataOpsTest, pod_id_to_node_name_test) {
   udf_tester.ForInput("123_uid").Expect("");
 }
 
+TEST_F(MetadataOpsTest, pod_id_to_host_ip_test) {
+  auto function_ctx = std::make_unique<FunctionContext>(metadata_state_, nullptr);
+  auto udf_tester = px::carnot::udf::UDFTester<PodIDToHostIPUDF>(std::move(function_ctx));
+  udf_tester.ForInput("1_uid").Expect("2.2.2.2");
+  // This pod is not available, should return empty.
+  udf_tester.ForInput("123_uid").Expect("");
+}
+
 TEST_F(MetadataOpsTest, pod_id_to_replicaset_name_test) {
   auto function_ctx = std::make_unique<FunctionContext>(metadata_state_, nullptr);
   auto udf_tester = px::carnot::udf::UDFTester<PodIDToReplicaSetNameUDF>(std::move(function_ctx));
@@ -651,6 +659,14 @@ TEST_F(MetadataOpsTest, pod_name_to_pod_ip) {
   udf_tester.ForInput("pl/nonexistent").Expect("");
 }
 
+TEST_F(MetadataOpsTest, pod_name_to_host_ip) {
+  auto function_ctx = std::make_unique<FunctionContext>(metadata_state_, nullptr);
+  auto udf_tester = px::carnot::udf::UDFTester<PodNameToHostIPUDF>(std::move(function_ctx));
+  udf_tester.ForInput("pl/running_pod").Expect("2.2.2.2");
+  udf_tester.ForInput("pl/terminating_pod").Expect("");
+  udf_tester.ForInput("pl/nonexistent").Expect("");
+}
+
 TEST_F(MetadataOpsTest, container_id_to_container_status) {
   ContainerIDToContainerStatusUDF status_udf;
 

src/shared/k8s/metadatapb/test_proto.h

@@ -56,6 +56,7 @@ qos_class: QOS_CLASS_GUARANTEED
 node_name: "test_node"
 hostname: "test_host"
 pod_ip: "1.1.1.1"
+host_ip: "2.2.2.2"
 phase: RUNNING
 message: "Running message"
 reason: "Running reason"

src/shared/metadata/k8s_objects.h

@@ -322,11 +322,13 @@ class PodInfo : public K8sMetadataObject {
   void set_node_name(std::string_view node_name) { node_name_ = node_name; }
   void set_hostname(std::string_view hostname) { hostname_ = hostname; }
   void set_pod_ip(std::string_view pod_ip) { pod_ip_ = pod_ip; }
+  void set_host_ip(std::string_view host_ip) { host_ip_ = host_ip; }
   void set_pod_labels(std::string labels) { labels_ = labels; }
 
   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& host_ip() const { return host_ip_; }
   const std::string& labels() const { return labels_; }
 
   const absl::flat_hash_set<std::string>& containers() const { return containers_; }
@@ -367,6 +369,8 @@ class PodInfo : public K8sMetadataObject {
   std::string node_name_;
   std::string hostname_;
   std::string pod_ip_;
+  // IP address of the host (node) the pod is assigned to.
+  std::string host_ip_;
   std::string labels_;
 };
 

src/shared/metadata/metadata_state.cc

@@ -287,6 +287,7 @@ Status K8sMetadataState::HandlePodUpdate(const PodUpdate& update) {
   pod_info->set_node_name(update.node_name());
   pod_info->set_hostname(update.hostname());
   pod_info->set_pod_ip(update.pod_ip());
+  pod_info->set_host_ip(update.host_ip());
   pod_info->set_phase(ConvertToPodPhase(update.phase()));
   pod_info->set_conditions(ConvertToPodConditions(update.conditions()));
   pod_info->set_phase_message(update.message());

src/shared/metadata/metadata_state_test.cc

@@ -385,6 +385,7 @@ TEST(K8sMetadataStateTest, HandlePodUpdate) {
   EXPECT_EQ("a_node", pod_info->node_name());
   EXPECT_EQ("a_host", pod_info->hostname());
   EXPECT_EQ("1.2.3.4", pod_info->pod_ip());
+  EXPECT_EQ("5.6.7.8", pod_info->host_ip());
   EXPECT_THAT(pod_info->containers(), UnorderedElementsAre("container0_uid"));
   EXPECT_THAT(pod_info->owner_references(),
               UnorderedElementsAre(OwnerReference{"1234", "rs0", "ReplicaSet"}));