From b8d2e8b988b3110e2e83fb5cd3f1587944b07b15 Mon Sep 17 00:00:00 2001
From: Xu <hxu398@krishna-srv1.ece.gatech.edu>
Date: Mon, 3 Feb 2025 23:10:21 -0500
Subject: [PATCH] Handle optional etfeeder node; introduce INVALID_COMM_TYPE
 enum

---
 schema/protobuf/et_def.proto  |  1 +
 src/feeder/et_feeder_node.cpp | 28 +++++++++++++---------------
 src/feeder/et_feeder_node.h   | 24 ++++++++++++++----------
 3 files changed, 28 insertions(+), 25 deletions(-)

diff --git a/schema/protobuf/et_def.proto b/schema/protobuf/et_def.proto
index 965cec86..1a2d3ea9 100644
--- a/schema/protobuf/et_def.proto
+++ b/schema/protobuf/et_def.proto
@@ -127,6 +127,7 @@ enum CollectiveCommType {
   REDUCE_SCATTER = 7;
   REDUCE_SCATTER_BLOCK = 8;
   BARRIER = 9;
+  INVALID_COMM_TYPE = 99;
 }
 
 message Node {
diff --git a/src/feeder/et_feeder_node.cpp b/src/feeder/et_feeder_node.cpp
index 4b5e99e5..fb08b5f0 100644
--- a/src/feeder/et_feeder_node.cpp
+++ b/src/feeder/et_feeder_node.cpp
@@ -17,6 +17,8 @@ ETFeederNode::ETFeederNode(std::shared_ptr<ChakraProtoMsg::Node> node) {
       this->is_cpu_op_ = static_cast<bool>(attr.bool_val());
     } else if (attr_name == "num_ops") {
       this->num_ops_ = static_cast<uint64_t>(attr.int64_val());
+    } else if (attr_name == "tensor_loc_") {
+      this->tensor_loc_ = attr.uint32_val();
     } else if (attr_name == "tensor_size") {
       this->tensor_size_ = attr.uint64_val();
     } else if (attr_name == "comm_type") {
@@ -106,41 +108,37 @@ uint64_t ETFeederNode::runtime() {
 }
 
 uint64_t ETFeederNode::num_ops() {
-  return num_ops_;
-}
-
-uint32_t ETFeederNode::tensor_loc() {
-  return tensor_loc_;
+  return num_ops_.value_or(0);
 }
 
 uint64_t ETFeederNode::tensor_size() {
-  return tensor_size_;
+  return tensor_size_.value_or(0);
 }
 
 ChakraProtoMsg::CollectiveCommType ETFeederNode::comm_type() {
-  return comm_type_;
+  return comm_type_.value_or(ChakraProtoMsg::CollectiveCommType::INVALID_COMM_TYPE);
 }
 
 uint32_t ETFeederNode::comm_priority() {
-  return comm_priority_;
+  return comm_priority_.value_or(0);
 }
 
 uint64_t ETFeederNode::comm_size() {
-  return comm_size_;
+  return comm_size_.value_or(0);
 }
 
 uint32_t ETFeederNode::comm_src() {
-  return comm_src_;
+  return comm_src_.value_or(0);
 }
 
 uint32_t ETFeederNode::comm_dst() {
-  return comm_dst_;
+  return comm_dst_.value_or(0);
 }
 
 uint32_t ETFeederNode::comm_tag() {
-  return comm_tag_;
+  return comm_tag_.value_or(0);
 }
 
-string ETFeederNode::pg_name() {
-  return pg_name_;
-}
+std::string ETFeederNode::pg_name() {
+  return pg_name_.value_or("");
+}
\ No newline at end of file
diff --git a/src/feeder/et_feeder_node.h b/src/feeder/et_feeder_node.h
index c0aede48..8e76fce3 100644
--- a/src/feeder/et_feeder_node.h
+++ b/src/feeder/et_feeder_node.h
@@ -4,6 +4,7 @@
 #include <unordered_map>
 #include <unordered_set>
 #include <vector>
+#include <optional>
 
 #include "et_def.pb.h"
 
@@ -53,20 +54,23 @@ class ETFeederNode {
   std::unordered_map<std::string, const ChakraProtoMsg::AttributeProto&>
       other_attrs_{};
 
+  // required metadata
   uint64_t id_;
   std::string name_;
   bool is_cpu_op_;
   uint64_t runtime_;
-  uint64_t num_ops_;
-  uint32_t tensor_loc_;
-  uint64_t tensor_size_;
-  ChakraProtoMsg::CollectiveCommType comm_type_;
-  uint32_t comm_priority_;
-  uint64_t comm_size_;
-  uint32_t comm_src_;
-  uint32_t comm_dst_;
-  uint32_t comm_tag_;
-  std::string pg_name_;
+
+  // optional metadata
+  std::optional<uint64_t> num_ops_;
+  std::optional<uint32_t> tensor_loc_;
+  std::optional<uint64_t> tensor_size_;
+  std::optional<ChakraProtoMsg::CollectiveCommType> comm_type_;
+  std::optional<uint32_t> comm_priority_;
+  std::optional<uint64_t> comm_size_;
+  std::optional<uint32_t> comm_src_;
+  std::optional<uint32_t> comm_dst_;
+  std::optional<uint32_t> comm_tag_;
+  std::optional<std::string> pg_name_;
 };
 
 } // namespace Chakra