Fix RayStageSpecKeys enum membership check for Python < 3.12 compatibility (#1039)

* Initial plan

* Fix RayStageSpecKeys enum membership check for Python < 3.12 compatibility

Co-authored-by: VibhuJawa <[email protected]>

* Add test for RayStageSpecKeys compatibility fix

Co-authored-by: VibhuJawa <[email protected]>

* Fix ruff linting issues in test file - remove trailing whitespace and add final newline

Co-authored-by: VibhuJawa <[email protected]>

* Address code review feedback - import RayStageSpecKeys instead of redefining it in test

Co-authored-by: VibhuJawa <[email protected]>

* Remove __main__ block from test file as requested in code review

Co-authored-by: VibhuJawa <[email protected]>

* Move RayStageSpecKeys compatibility test to proper location and clean up format

Co-authored-by: VibhuJawa <[email protected]>

---------

Co-authored-by: copilot-swe-agent[bot] <[email protected]>
Co-authored-by: VibhuJawa <[email protected]>
Co-authored-by: Vibhu Jawa <[email protected]>

Author: Copilot

nemo_curator/backends/experimental/ray_data/adapter.py

@@ -46,7 +46,7 @@ def __init__(self, stage: ProcessingStage):
 
         # Go through all the keys in the ray_stage_spec and raise error if they are not in RayStageSpecKeys
         for key in self.stage.ray_stage_spec():
-            if key not in RayStageSpecKeys:
+            if key not in {e.value for e in RayStageSpecKeys}:
                 msg = f"Invalid key {key} in ray_stage_spec for stage {self.stage}"
                 raise ValueError(msg)
 

tests/backends/experimental/test_utils.py

@@ -20,7 +20,7 @@
 from pytest import LogCaptureFixture
 
 from nemo_curator.backends.base import NodeInfo, WorkerMetadata
-from nemo_curator.backends.experimental.utils import execute_setup_on_node
+from nemo_curator.backends.experimental.utils import RayStageSpecKeys, execute_setup_on_node
 from nemo_curator.stages.base import ProcessingStage
 from nemo_curator.stages.resources import Resources
 
@@ -93,3 +93,26 @@ def setup_on_node(
         assert len(matching_logs) == len(ray.nodes()), (
             f"Expected {len(ray.nodes())} logs for setup on node for 2 stages, got {len(matching_logs)}: {matching_logs}"
         )
+
+
+class TestRayStageSpecKeys:
+    """Test class for RayStageSpecKeys enum compatibility."""
+
+    def test_enum_membership_compatibility(self):
+        """Test that the fixed pattern works across Python versions."""
+        # Test data
+        valid_keys = ["is_actor_stage", "is_fanout_stage", "is_lsh_stage"]
+        invalid_keys = ["invalid_key", "another_bad_key"]
+
+        # Test the fixed pattern - this is what's now used in the adapter
+        enum_values = {e.value for e in RayStageSpecKeys}
+
+        # Testing valid keys
+        for key in valid_keys:
+            result = key not in enum_values
+            assert result is False, f"Valid key '{key}' should be found in enum values"
+
+        # Testing invalid keys
+        for key in invalid_keys:
+            result = key not in enum_values
+            assert result is True, f"Invalid key '{key}' should not be found in enum values"