Skip RES HBM tests when op absent in OSS wheel (pytorch#6233)

Summary:
X-link: https://github.com/facebookresearch/FBGEMM/pull/3120


The RES (raw embedding streaming) HBM-lane tests in res_enabled_tables_test.py
reach torch.ops.fbgemm.masked_index_select, which is registered in the SSD
split-embeddings-cache extension (src/ssd_split_embeddings_cache/). That
directory is compiled only by the internal Buck build; the OSS CMake wheel
excludes it entirely, so the op is absent in OSS. Every HBM-lane test therefore
dies on an _OpNamespace lookup ("'fbgemm' object has no attribute
'masked_index_select'"), plus cascading _res_* attribute errors, turning the
OSS 'main' GPU CI red (30 failed / 13 passed).

Guard the single chokepoint all HBM-lane tests construct through
(_build_mixed_tbe, reached directly and via _drain_tbe/_device_tbe) with a
SkipTest when masked_index_select is not registered in the running wheel. On any
build that has the op (internal, and OSS once the extension is shipped) the guard
is a no-op and all tests run unchanged. The non-HBM allowlist tests use
_build_tbe, never touch the op, and continue to run everywhere.

This is the fast unblock for OSS trunk/nightly. The proper long-term fix is to
add src/ssd_split_embeddings_cache/ to the OSS CMake build so the op is present.

Reviewed By: FriedCosey

Differential Revision: D117750928

Author: gchalump

fbgemm_gpu/test/tbe/training/res_enabled_tables_test.py

@@ -33,6 +33,19 @@
 # arithmetic against. One name because the two have to agree.
 ROWS = 64
 
+# The RES HBM streaming lane reaches ``torch.ops.fbgemm.masked_index_select``,
+# which is registered in the SSD split-embeddings-cache extension
+# (``src/ssd_split_embeddings_cache/``). That directory is compiled only by the
+# internal Buck build; the OSS CMake wheel excludes it entirely, so the op is
+# absent in OSS and every HBM-lane test dies on an ``_OpNamespace`` lookup. Skip
+# those tests when the op isn't built into the running wheel. The non-HBM
+# allowlist tests (which never touch the op) still run everywhere.
+res_hbm_streaming_unavailable: tuple[bool, str] = (
+    not hasattr(torch.ops.fbgemm, "masked_index_select"),
+    "fbgemm.masked_index_select is not built into this wheel "
+    "(OSS CMake excludes src/ssd_split_embeddings_cache/)",
+)
+
 
 class ResEnabledTablesTest(unittest.TestCase):
     """
@@ -123,6 +136,8 @@ def _build_mixed_tbe(
         res_hbm_drain_interval: int = 1,
     ) -> SplitTableBatchedEmbeddingBagsCodegen:
         """One table per name, each with its own placement, row count and dim."""
+        if res_hbm_streaming_unavailable[0]:
+            raise unittest.SkipTest(res_hbm_streaming_unavailable[1])
         n = len(table_names)
         rows = heights if heights is not None else [ROWS] * n
         widths = dims if dims is not None else [16] * n