Silent-failure ODS counters (#5850)

Summary:

X-link: facebookresearch/FBGEMM#2768

Add OBC counters for 3 silent failure modes in the Raw Embedding Streaming (RES) path, closing the must-have bar of T273802725 (subtask of T269497764, RES code-side observability this half). Today these failure sites only emit `XLOG(ERR)` / `LOG(ERROR)` — nothing aggregates, so on-call cannot trigger off them or query them in ODS without `tail -f` on individual trainer hosts.

Counters added (all via `TrainingPsOdsLogger::bumpKey`, OBC category `raw_embedding_streaming`):

| Counter | Site |
|---|---|
| `res.fail.shard_size_mismatch` | `raw_embedding_streamer.cpp` — both the outer `indices.size(0) != weights.size(0)` guard and the per-shard `weights_masked.size(0) != rows_in_shard` guard in `tensor_stream()` |
| `res.fail.set_embeddings_rpc` | `raw_embedding_streamer.cpp` — wraps `co_await res_client->co_setEmbeddings(req)` in try/catch; bumps then re-throws to preserve trainer-thread termination |
| `res.fail.write_sparse_throw` | `TrainingPsHandler.cpp` `catch (const std::exception& e)` block at the `writeSparseListOfTensors` site in `co_stream_tensors` |

Backend (addresses review: use OBC, not fb303): all three route through the existing `TrainingPsOdsLogger` OBC client. The handler already owned an `ods_logger_`; the fbgemm-side `RawEmbeddingStreamer` now holds a `TrainingPsOdsLogger` member (forward-declared in the OSS-mirrored header, constructed only when streaming is enabled). OBC reaches ODS through the host-level agent with no per-process fb303 export/scrape config, so this revision deletes the fb303 `ResFailCounters.h` helper (and its `call_once` / `addStatExportType` / `FOLLY_EXPORT` machinery) entirely.

Counters bump AT the failure site (inside catch / right before `continue`) so they reflect what actually happened — applying the correctness lesson from D104076551 V4 (don't preemptively count attempts at the top of the function). Per-TBE-per-rank granularity is fine: OBC aggregates per (entity, key), so per-TBE bumps land on the same per-host key rather than separate rows.

Existing log lines preserved (additive). Existing throw-on-failure semantics preserved — the `co_setEmbeddings` catch re-throws after bumping so trainer-thread termination is unchanged.

Out of scope (intentional, per T269497764 boundary):
- MPSC `weights_to_stream_queue_` near-full counter — queue is unbounded; deferred to the queue-depth gauge work in T273802756.
- Detectors / alerts on these counters — downstream consumer work.

Reviewed By: FriedCosey

Differential Revision: D107811590

Author: adityas-meta

fbgemm_gpu/include/fbgemm_gpu/split_embeddings_cache/raw_embedding_streamer.h

@@ -14,6 +14,12 @@
 
 #include <utility>
 
+#ifdef FBGEMM_FBCODE
+namespace facebook::aiplatform::gmpp::experimental::training_ps {
+class TrainingPsOdsLogger;
+} // namespace facebook::aiplatform::gmpp::experimental::training_ps
+#endif
+
 namespace fbgemm_gpu {
 
 struct StreamQueueItem {
@@ -125,6 +131,12 @@ class RawEmbeddingStreamer : public torch::jit::CustomClassHolder {
   std::unique_ptr<std::thread> weights_stream_thread_;
   folly::UMPSCQueue<StreamQueueItem, true> weights_to_stream_queue_;
   std::unique_ptr<std::thread> stream_tensor_copy_thread_;
+  // OBC logger for RES silent-failure counters (res.fail.*). Emits to the
+  // host-level OBC agent, so it reaches ODS from the trainer process without
+  // per-process fb303 scrape config. Only constructed when streaming is on.
+  std::unique_ptr<facebook::aiplatform::gmpp::experimental::training_ps::
+                      TrainingPsOdsLogger>
+      ods_logger_;
 #endif
 };
 

fbgemm_gpu/src/split_embeddings_cache/raw_embedding_streamer.cpp

@@ -10,6 +10,7 @@
 #include <folly/coro/BlockingWait.h>
 #include <folly/stop_watch.h>
 #include <utility>
+#include "aiplatform/gmpp/experimental/training_ps/TrainingPsOdsLogger.h"
 #include "aiplatform/gmpp/experimental/training_ps/gen-cpp2/TrainingParameterServerService.h"
 #include "caffe2/torch/fb/distributed/wireSerializer/WireSerializer.h"
 #include "servicerouter/client/cpp2/ClientParams.h"
@@ -49,6 +50,7 @@ get_res_client(int64_t res_server_port) {
                                  TrainingParameterServerService>>(
       "realtime.delta.publish.esr", params);
 }
+
 #endif
 
 /// Read a scalar value from a tensor that is maybe a UVM tensor
@@ -167,6 +169,9 @@ RawEmbeddingStreamer::RawEmbeddingStreamer(
     // The first call to get the client is expensive, so eagerly get it here
     auto _eager_client = get_res_client(res_server_port_);
 
+    ods_logger_ = std::make_unique<facebook::aiplatform::gmpp::experimental::
+                                       training_ps::TrainingPsOdsLogger>();
+
     weights_stream_thread_ = std::make_unique<std::thread>([this] {
       while (!stop_) {
         auto stream_item_ptr = weights_to_stream_queue_.try_peek();
@@ -319,6 +324,9 @@ folly::coro::Task<void> RawEmbeddingStreamer::tensor_stream(
     XLOG(ERR) << "[TBE_ID" << unique_id_
               << "] Indices and weights size mismatched " << indices.size(0)
               << " " << weights.size(0);
+    if (ods_logger_) {
+      ods_logger_->bumpKey("shard_size_mismatch", 1);
+    }
     co_return;
   }
   folly::stop_watch<std::chrono::milliseconds> stop_watch;
@@ -399,6 +407,9 @@ folly::coro::Task<void> RawEmbeddingStreamer::tensor_stream(
           << "] don't send the request for size mismatched tensors table: "
           << rows_in_shard << " weights: " << weights_masked.size(0)
           << " global_indices: " << global_indices_masked.numel();
+      if (ods_logger_) {
+        ods_logger_->bumpKey("shard_size_mismatch", 1);
+      }
       continue;
     }
     SetEmbeddingsRequest req;
@@ -420,7 +431,17 @@ folly::coro::Task<void> RawEmbeddingStreamer::tensor_stream(
       req.runtimeMeta() =
           torch::distributed::wireDumpTensor(runtime_meta_masked);
     }
-    co_await res_client->co_setEmbeddings(req);
+    try {
+      co_await res_client->co_setEmbeddings(req);
+    } catch (const std::exception& e) {
+      if (ods_logger_) {
+        ods_logger_->bumpKey("set_embeddings_rpc", 1);
+      }
+      XLOG(ERR) << "[TBE_ID" << unique_id_
+                << "] co_setEmbeddings threw on shard " << i << ": "
+                << e.what();
+      throw;
+    }
   }
   co_return;
 }

fbgemm_gpu/src/split_embeddings_cache/tests/raw_embedding_streamer_test.cpp

@@ -237,6 +237,61 @@ TEST(RawEmbeddingStreamerTest, TestStreamE2E) {
   streamer->join_weights_stream_thread();
 }
 
+TEST(RawEmbeddingStreamerTest, TestCoSetEmbeddingsThrowPropagates) {
+  std::vector<std::string> table_names = {"tb1", "tb2", "tb3"};
+  std::vector<int64_t> table_offsets = {0, 100, 300};
+  std::vector<int64_t> table_sizes = {0, 50, 200, 300};
+
+  auto streamer = getRawEmbeddingStreamer(
+      "test_set_embeddings_throw",
+      true,
+      table_names,
+      table_offsets,
+      table_sizes);
+
+  auto mock_service = std::make_shared<MockTrainingParameterServerService>();
+  auto mock_server =
+      std::make_shared<apache::thrift::ScopedServerInterfaceThread>(
+          mock_service,
+          "::1",
+          0,
+          facebook::services::TLSConfig::applyDefaultsToThriftServer);
+  auto& mock_client_factory =
+      facebook::servicerouter::getMockSRClientFactory(false /* strict */);
+  mock_client_factory.registerMockService(
+      "realtime.delta.publish.esr", mock_server);
+
+  EXPECT_CALL(*mock_service, co_setEmbeddings(_))
+      .WillRepeatedly(
+          folly::coro::gmock_helpers::CoInvoke(
+              [](std::unique_ptr<aiplatform::gmpp::experimental::training_ps::
+                                     SetEmbeddingsRequest>)
+                  -> folly::coro::Task<
+                      std::unique_ptr<aiplatform::gmpp::experimental::
+                                          training_ps::SetEmbeddingsResponse>> {
+                throw std::runtime_error("simulated RPC failure");
+                co_return std::make_unique<
+                    aiplatform::gmpp::experimental::training_ps::
+                        SetEmbeddingsResponse>();
+              }));
+
+  auto indices = at::tensor(
+      {10, 2, 1, 150, 170, 230, 280},
+      at::TensorOptions().device(at::kCPU).dtype(at::kLong));
+  auto weights = at::randn(
+      {indices.size(0), EMBEDDING_DIMENSION},
+      at::TensorOptions().device(at::kCPU).dtype(c10::kFloat));
+
+  // Thrift repackages server-side exceptions into TApplicationException, so
+  // assert that the exception still propagates (preserved contract). The
+  // catch block bumps set_embeddings_rpc via the OBC logger before
+  // rethrowing; OBC counters are not in-process readable, so the bump itself
+  // is not unit-asserted here (the logger has no in-test sink).
+  EXPECT_ANY_THROW(
+      folly::coro::blockingWait(streamer->tensor_stream(
+          indices, weights, std::nullopt, std::nullopt)));
+}
+
 TEST(RawEmbeddingStreamerTest, TestMismatchedIndicesWeights) {
   std::vector<std::string> table_names = {"tb1", "tb2", "tb3"};
   std::vector<int64_t> table_offsets = {0, 100, 300};
@@ -258,7 +313,10 @@ TEST(RawEmbeddingStreamerTest, TestMismatchedIndicesWeights) {
   mock_client_factory.registerMockService(
       "realtime.delta.publish.esr", mock_server);
 
-  // Test with mismatched sizes - should not call service
+  // Test with mismatched sizes - should not call service. The
+  // shard_size_mismatch counter is bumped in this path, but OBC
+  // counters are not in-process readable, so we verify the behavior (no RPC)
+  // rather than the counter value.
   auto indices = at::tensor(
       {10, 2, 1}, at::TensorOptions().device(at::kCPU).dtype(at::kLong));
   auto weights = at::randn(