Add UDF to discover namespace where Vizier is deployed (#1307)

rosenhouse · web-flow · commit 6182360d78f6 · 2023-05-11T11:51:28.000-07:00
Summary: Adds a new `px.vizier_namespace()` function usable from PxL scripts to discover the namespace that the Vizier components are deployed into. Type of change: /kind feature Test Plan: Added some unit tests. Manually tested on a minikube. Relevant Issues: Fixes #1304 Changelog Message: ```release-note Add `px.vizier_namespace()` UDF to get namespace that Vizier is running in. ``` --------- Signed-off-by: Gabe Rosenhouse <grosenhouse@vmware.com>
diff --git a/k8s/vizier/base/kelvin_deployment.yaml b/k8s/vizier/base/kelvin_deployment.yaml
@@ -95,7 +95,7 @@ spec:
               key: cluster-name
               name: pl-cluster-secrets
               optional: true
-        - name: PL_NAMESPACE
+        - name: PL_POD_NAMESPACE
           valueFrom:
             fieldRef:
               fieldPath: metadata.namespace
diff --git a/k8s/vizier/pem/base/pem_daemonset.yaml b/k8s/vizier/pem/base/pem_daemonset.yaml
@@ -91,6 +91,10 @@ spec:
           valueFrom:
             fieldRef:
               fieldPath: metadata.name
+        - name: PL_POD_NAMESPACE
+          valueFrom:
+            fieldRef:
+              fieldPath: metadata.namespace
         - name: PL_HOST_IP
           valueFrom:
             fieldRef:
diff --git a/src/carnot/funcs/metadata/metadata_ops.cc b/src/carnot/funcs/metadata/metadata_ops.cc
@@ -128,6 +128,7 @@ void RegisterMetadataOpsOrDie(px::carnot::udf::Registry* registry) {
   registry->RegisterOrDie<HostNumCPUsUDF>("_exec_host_num_cpus");
   registry->RegisterOrDie<VizierIDUDF>("vizier_id");
   registry->RegisterOrDie<VizierNameUDF>("vizier_name");
+  registry->RegisterOrDie<VizierNamespaceUDF>("vizier_namespace");
   registry->RegisterOrDie<GetClusterCIDRRangeUDF>("get_cidrs");
 
   /*****************************************
diff --git a/src/carnot/funcs/metadata/metadata_ops.h b/src/carnot/funcs/metadata/metadata_ops.h
@@ -3120,6 +3120,21 @@ class VizierNameUDF : public ScalarUDF {
   }
 };
 
+class VizierNamespaceUDF : public ScalarUDF {
+ public:
+  StringValue Exec(FunctionContext* ctx) {
+    auto md = GetMetadataState(ctx);
+    return md->vizier_namespace();
+  }
+
+  static udf::ScalarUDFDocBuilder Doc() {
+    return udf::ScalarUDFDocBuilder("Get the namespace where Vizier is deployed.")
+        .Details("Get the Kubernetes namespace in which the Vizier is deployed.")
+        .Example("df.vizier_namespace = px.vizier_namespace()")
+        .Returns("The Kubernetes namespace in which Vizier is deployed");
+  }
+};
+
 class CreateUPIDUDF : public udf::ScalarUDF {
  public:
   UInt128Value Exec(FunctionContext* ctx, Int64Value pid, Int64Value pid_start_time) {
diff --git a/src/carnot/funcs/metadata/metadata_ops_test.cc b/src/carnot/funcs/metadata/metadata_ops_test.cc
@@ -52,7 +52,8 @@ class MetadataOpsTest : public ::testing::Test {
     vizier_id_ = sole::uuid4();
     metadata_state_ = std::make_shared<px::md::AgentMetadataState>(
         /* hostname */ "myhost",
-        /* asid */ 1, /* pid */ 123, agent_id_, "mypod", vizier_id_, "myvizier");
+        /* asid */ 1, /* pid */ 123, agent_id_, "mypod", vizier_id_, "myvizier",
+        "myviziernamespace");
     // Apply updates to metadata state.
     updates_ =
         std::make_unique<moodycamel::BlockingConcurrentQueue<std::unique_ptr<ResourceUpdate>>>();
@@ -1269,7 +1270,7 @@ TEST_F(MetadataOpsTest, in_value_or_array_test) {
 TEST_F(MetadataOpsTest, vizier_id_test) {
   auto metadata_state = std::make_shared<px::md::AgentMetadataState>(
       /* hostname */ "myhost",
-      /* asid */ 1, /* pid */ 123, agent_id_, "mypod", vizier_id_, "myvizier");
+      /* asid */ 1, /* pid */ 123, agent_id_, "mypod", vizier_id_, "myvizier", "myviziernamespace");
   auto function_ctx = std::make_unique<FunctionContext>(metadata_state, nullptr);
   auto udf_tester = px::carnot::udf::UDFTester<VizierIDUDF>(std::move(function_ctx));
   udf_tester.ForInput().Expect(vizier_id_.str());
@@ -1278,7 +1279,8 @@ TEST_F(MetadataOpsTest, vizier_id_test) {
 TEST_F(MetadataOpsTest, empty_vizier_id_test) {
   auto metadata_state = std::make_shared<px::md::AgentMetadataState>(
       /* hostname */ "myhost",
-      /* asid */ 1, /* pid */ 123, agent_id_, "mypod", sole::uuid(), "myvizier");
+      /* asid */ 1, /* pid */ 123, agent_id_, "mypod", sole::uuid(), "myvizier",
+      "myviziernamespace");
   auto function_ctx = std::make_unique<FunctionContext>(metadata_state, nullptr);
   auto udf_tester = px::carnot::udf::UDFTester<VizierIDUDF>(std::move(function_ctx));
   udf_tester.ForInput().Expect("00000000-0000-0000-0000-000000000000");
@@ -1287,16 +1289,25 @@ TEST_F(MetadataOpsTest, empty_vizier_id_test) {
 TEST_F(MetadataOpsTest, vizier_name_test) {
   auto metadata_state = std::make_shared<px::md::AgentMetadataState>(
       /* hostname */ "myhost",
-      /* asid */ 1, /* pid */ 123, agent_id_, "mypod", vizier_id_, "myvizier");
+      /* asid */ 1, /* pid */ 123, agent_id_, "mypod", vizier_id_, "myvizier", "myviziernamespace");
   auto function_ctx = std::make_unique<FunctionContext>(metadata_state, nullptr);
   auto udf_tester = px::carnot::udf::UDFTester<VizierNameUDF>(std::move(function_ctx));
   udf_tester.ForInput().Expect("myvizier");
 }
 
+TEST_F(MetadataOpsTest, vizier_namespace_test) {
+  auto metadata_state = std::make_shared<px::md::AgentMetadataState>(
+      /* hostname */ "myhost",
+      /* asid */ 1, /* pid */ 123, agent_id_, "mypod", vizier_id_, "myvizier", "myviziernamespace");
+  auto function_ctx = std::make_unique<FunctionContext>(metadata_state, nullptr);
+  auto udf_tester = px::carnot::udf::UDFTester<VizierNamespaceUDF>(std::move(function_ctx));
+  udf_tester.ForInput().Expect("myviziernamespace");
+}
+
 TEST_F(MetadataOpsTest, basic_upid_test) {
   auto metadata_state = std::make_shared<px::md::AgentMetadataState>(
       /* hostname */ "myhost",
-      /* asid */ 1, /* pid */ 123, agent_id_, "mypod", vizier_id_, "myvizier");
+      /* asid */ 1, /* pid */ 123, agent_id_, "mypod", vizier_id_, "myvizier", "myviziernamespace");
   auto function_ctx = std::make_unique<FunctionContext>(metadata_state, nullptr);
   auto udf_tester = px::carnot::udf::UDFTester<CreateUPIDUDF>(std::move(function_ctx));
   udf_tester.ForInput(100, 23123).Expect(md::UPID(1, 100, 23123).value());
@@ -1305,7 +1316,7 @@ TEST_F(MetadataOpsTest, basic_upid_test) {
 TEST_F(MetadataOpsTest, basic_upid_with_asid_test) {
   auto metadata_state = std::make_shared<px::md::AgentMetadataState>(
       /* hostname */ "myhost",
-      /* asid */ 1, /* pid */ 123, agent_id_, "mypod", vizier_id_, "myvizier");
+      /* asid */ 1, /* pid */ 123, agent_id_, "mypod", vizier_id_, "myvizier", "myviziernamespace");
   auto function_ctx = std::make_unique<FunctionContext>(metadata_state, nullptr);
   auto udf_tester = px::carnot::udf::UDFTester<CreateUPIDWithASIDUDF>(std::move(function_ctx));
   udf_tester.ForInput(2, 100, 23123).Expect(md::UPID(2, 100, 23123).value());
@@ -1314,7 +1325,7 @@ TEST_F(MetadataOpsTest, basic_upid_with_asid_test) {
 TEST_F(MetadataOpsTest, get_cidrs) {
   auto metadata_state = std::make_shared<px::md::AgentMetadataState>(
       /* hostname */ "myhost",
-      /* asid */ 1, /* pid */ 123, agent_id_, "mypod", vizier_id_, "myvizier");
+      /* asid */ 1, /* pid */ 123, agent_id_, "mypod", vizier_id_, "myvizier", "myviziernamespace");
   px::CIDRBlock pod_cidr;
   std::string pod_cidr_str("10.0.0.20/32");
   ASSERT_OK(px::ParseCIDRBlock(pod_cidr_str, &pod_cidr));
diff --git a/src/shared/metadata/metadata_state.cc b/src/shared/metadata/metadata_state.cc
@@ -513,7 +513,7 @@ Status K8sMetadataState::CleanupExpiredMetadata(int64_t retention_time_ns) {
 
 std::shared_ptr<AgentMetadataState> AgentMetadataState::CloneToShared() const {
   auto state = std::make_shared<AgentMetadataState>(hostname_, asid_, pid_, agent_id_, pod_name_,
-                                                    vizier_id_, vizier_name_);
+                                                    vizier_id_, vizier_name_, vizier_namespace_);
   state->last_update_ts_ns_ = last_update_ts_ns_;
   state->epoch_id_ = epoch_id_;
   state->k8s_metadata_state_ = k8s_metadata_state_->Clone();
diff --git a/src/shared/metadata/metadata_state.h b/src/shared/metadata/metadata_state.h
@@ -317,17 +317,19 @@ class AgentMetadataState : NotCopyable {
   AgentMetadataState() = delete;
   explicit AgentMetadataState(uint32_t asid, uint32_t pid)
       : AgentMetadataState(/* hostname */ "unknown", asid, pid, sole::uuid(),
-                           /* pod_name */ "unknown", sole::uuid(), "") {}
+                           /* pod_name */ "unknown", sole::uuid(), "", "") {}
 
   AgentMetadataState(std::string_view hostname, uint32_t asid, uint32_t pid, AgentID agent_id,
-                     std::string_view pod_name, sole::uuid vizier_id, std::string_view vizier_name)
+                     std::string_view pod_name, sole::uuid vizier_id, std::string_view vizier_name,
+                     std::string_view vizier_namespace)
       : hostname_(std::string(hostname)),
         pod_name_(std::string(pod_name)),
         asid_(asid),
         pid_(pid),
         agent_id_(agent_id),
         vizier_id_(vizier_id),
         vizier_name_(std::string(vizier_name)),
+        vizier_namespace_(std::string(vizier_namespace)),
         k8s_metadata_state_(new K8sMetadataState()) {}
 
   const std::string& hostname() const { return hostname_; }
@@ -338,6 +340,7 @@ class AgentMetadataState : NotCopyable {
 
   const sole::uuid& vizier_id() const { return vizier_id_; }
   const std::string& vizier_name() const { return vizier_name_; }
+  const std::string& vizier_namespace() const { return vizier_namespace_; }
 
   int64_t last_update_ts_ns() const { return last_update_ts_ns_; }
   void set_last_update_ts_ns(int64_t last_update_ts_ns) { last_update_ts_ns_ = last_update_ts_ns; }
@@ -405,6 +408,7 @@ class AgentMetadataState : NotCopyable {
 
   sole::uuid vizier_id_;
   std::string vizier_name_;
+  std::string vizier_namespace_;
 
   std::unique_ptr<K8sMetadataState> k8s_metadata_state_;
 
diff --git a/src/shared/metadata/standalone_state_manager.h b/src/shared/metadata/standalone_state_manager.h
@@ -38,7 +38,7 @@ class StandaloneAgentMetadataStateManager : public AgentMetadataStateManager {
                                       sole::uuid agent_id) {
     agent_metadata_state_ =
         std::make_shared<AgentMetadataState>(hostname, asid, pid, agent_id,
-                                             /*pod_name=*/"", sole::uuid(), "standalone_pem");
+                                             /*pod_name=*/"", sole::uuid(), "standalone_pem", "");
   }
   virtual ~StandaloneAgentMetadataStateManager() = default;
   AgentMetadataFilter* metadata_filter() const override { return nullptr; }
diff --git a/src/shared/metadata/state_manager.h b/src/shared/metadata/state_manager.h
@@ -122,11 +122,11 @@ class AgentMetadataStateManagerImpl : public AgentMetadataStateManager {
                                 std::string pod_name, sole::uuid agent_id, bool collects_data,
                                 const px::system::Config& config,
                                 AgentMetadataFilter* metadata_filter, sole::uuid vizier_id,
-                                std::string vizier_name)
+                                std::string vizier_name, std::string vizier_namespace)
       : pod_name_(pod_name), collects_data_(collects_data), metadata_filter_(metadata_filter) {
     md_reader_ = std::make_unique<CGroupMetadataReader>(config);
-    agent_metadata_state_ = std::make_shared<AgentMetadataState>(hostname, asid, pid, agent_id,
-                                                                 pod_name, vizier_id, vizier_name);
+    agent_metadata_state_ = std::make_shared<AgentMetadataState>(
+        hostname, asid, pid, agent_id, pod_name, vizier_id, vizier_name, vizier_namespace);
   }
 
   AgentMetadataFilter* metadata_filter() const override { return metadata_filter_; }
diff --git a/src/shared/metadata/state_manager_test.cc b/src/shared/metadata/state_manager_test.cc
@@ -247,11 +247,13 @@ class AgentMetadataStateTest : public ::testing::Test {
   static constexpr char kHostname[] = "myhost";
   static constexpr char kPodName[] = "mypod";
   static constexpr char kVizierName[] = "myvizier";
+  static constexpr char kVizierNamespace[] = "myviziernamespace";
 
   AgentMetadataStateTest()
       : agent_id_(sole::uuid4()),
         vizier_id_(sole::uuid4()),
-        metadata_state_(kHostname, kASID, kPID, agent_id_, kPodName, vizier_id_, kVizierName) {}
+        metadata_state_(kHostname, kASID, kPID, agent_id_, kPodName, vizier_id_, kVizierName,
+                        kVizierNamespace) {}
 
   sole::uuid agent_id_;
   sole::uuid vizier_id_;
@@ -273,6 +275,7 @@ TEST_F(AgentMetadataStateTest, initialize_md_state) {
   EXPECT_EQ(agent_id_.str(), metadata_state_.agent_id().str());
   EXPECT_EQ(vizier_id_.str(), metadata_state_.vizier_id().str());
   EXPECT_EQ("myvizier", metadata_state_.vizier_name());
+  EXPECT_EQ("myviziernamespace", metadata_state_.vizier_namespace());
 
   K8sMetadataState* state = metadata_state_.k8s_metadata_state();
   EXPECT_THAT(state->pods_by_name(), UnorderedElementsAre(Pair(Pair("pl", "pod1"), "pod_id1")));
@@ -406,7 +409,8 @@ TEST_F(AgentMetadataStateTest, cidr_test) {
   AgentMetadataStateManagerImpl mgr("test_host", /*asid*/ 0, /*pid*/ 987, "test_pod",
                                     /*id*/ sole::uuid4(),
                                     /*collects_data*/ true, px::system::Config::GetInstance(),
-                                    &md_filter_, /*vizier_id*/ sole::uuid4(), "test_vizier");
+                                    &md_filter_, /*vizier_id*/ sole::uuid4(), "test_vizier",
+                                    "test_vizier_namespace");
 
   EXPECT_OK(mgr.PerformMetadataStateUpdate());
   // Should not be updated yet.
diff --git a/src/vizier/services/agent/kelvin/kelvin_main.cc b/src/vizier/services/agent/kelvin/kelvin_main.cc
@@ -38,7 +38,7 @@ DEFINE_string(mds_addr, gflags::StringFromEnv("PL_MDS_SVC_NAME", "vizier-metadat
 DEFINE_string(mds_port, gflags::StringFromEnv("PL_MDS_PORT", "50400"),
               "The port of the metadata service");
 
-DEFINE_string(namespace, gflags::StringFromEnv("PL_NAMESPACE", "pl"),
+DEFINE_string(namespace, gflags::StringFromEnv("PL_POD_NAMESPACE", "pl"),
               "The namespace of the Vizier instance");
 
 DEFINE_string(pod_ip, gflags::StringFromEnv("PL_POD_IP", ""),
diff --git a/src/vizier/services/agent/shared/manager/manager.cc b/src/vizier/services/agent/shared/manager/manager.cc
@@ -65,6 +65,9 @@ DEFINE_string(vizier_id, gflags::StringFromEnv("PL_VIZIER_ID", ""), "The ID of t
 DEFINE_string(vizier_name, gflags::StringFromEnv("PL_VIZIER_NAME", ""),
               "The name of the cluster according to vizier.");
 
+DEFINE_string(vizier_namespace, gflags::StringFromEnv("PL_POD_NAMESPACE", ""),
+              "The namespace in which vizier is deployed.");
+
 namespace px {
 namespace vizier {
 namespace agent {
@@ -285,7 +288,8 @@ Status Manager::PostRegisterHook(uint32_t asid) {
   mds_manager_ = std::make_unique<px::md::AgentMetadataStateManagerImpl>(
       info_.hostname, info_.asid, info_.pid, info_.pod_name, info_.agent_id,
       info_.capabilities.collects_data(), px::system::Config::GetInstance(),
-      agent_metadata_filter_.get(), sole::rebuild(FLAGS_vizier_id), FLAGS_vizier_name);
+      agent_metadata_filter_.get(), sole::rebuild(FLAGS_vizier_id), FLAGS_vizier_name,
+      FLAGS_vizier_namespace);
   // Register the Carnot callback for metadata.
   carnot_->RegisterAgentMetadataCallback(
       std::bind(&px::md::AgentMetadataStateManager::CurrentAgentMetadataState, mds_manager_.get()));
diff --git a/src/vizier/services/agent/shared/manager/test_utils.h b/src/vizier/services/agent/shared/manager/test_utils.h
@@ -57,7 +57,7 @@ class FakeAgentMetadataStateManager : public md::AgentMetadataStateManager {
   explicit FakeAgentMetadataStateManager(md::AgentMetadataFilter* metadata_filter)
       : metadata_filter_(metadata_filter) {
     metadata_state_ = std::make_shared<px::md::AgentMetadataState>(
-        "myhost", 1, 963, sole::uuid4(), "mypod", sole::uuid4(), "myvizier");
+        "myhost", 1, 963, sole::uuid4(), "mypod", sole::uuid4(), "myvizier", "myviziernamespace");
   }
 
   virtual ~FakeAgentMetadataStateManager() = default;

Original file line number	Diff line number	Diff line change
`@@ -38,7 +38,7 @@ class StandaloneAgentMetadataStateManager : public AgentMetadataStateManager {`
`38`	`38`	`sole::uuid agent_id) {`
`39`	`39`	`agent_metadata_state_ =`
`40`	`40`	`std::make_shared<AgentMetadataState>(hostname, asid, pid, agent_id,`
`41`		`- /pod_name=/"", sole::uuid(), "standalone_pem");`
	`41`	`+ /pod_name=/"", sole::uuid(), "standalone_pem", "");`
`42`	`42`	`}`
`43`	`43`	`virtual ~StandaloneAgentMetadataStateManager() = default;`
`44`	`44`	`AgentMetadataFilter* metadata_filter() const override { return nullptr; }`