From 2b3791ce0ef648068045511bc5c3d102d7d99281 Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Thu, 22 Jan 2026 19:15:06 +0000 Subject: [PATCH 1/6] Add validation checks for custom storage main --- .../src/validation_check/config_validator.py | 86 ++++ ...nd_resource_config_compatibility_checks.py | 167 ++++++++ .../libs/resource_config_checks.py | 51 +++ ...source_config_compatibility_checks_test.py | 403 ++++++++++++++++++ .../lib/resource_config_checks_test.py | 90 +++- 5 files changed, 796 insertions(+), 1 deletion(-) create mode 100644 python/gigl/src/validation_check/libs/gbml_and_resource_config_compatibility_checks.py create mode 100644 python/tests/unit/src/validation/lib/gbml_and_resource_config_compatibility_checks_test.py diff --git a/python/gigl/src/validation_check/config_validator.py b/python/gigl/src/validation_check/config_validator.py index 49af4637f..ec0ca4caf 100644 --- a/python/gigl/src/validation_check/config_validator.py +++ b/python/gigl/src/validation_check/config_validator.py @@ -15,6 +15,10 @@ assert_subgraph_sampler_output_exists, assert_trained_model_exists, ) +from gigl.src.validation_check.libs.gbml_and_resource_config_compatibility_checks import ( + check_inferencer_graph_store_compatibility, + check_trainer_graph_store_compatibility, +) from gigl.src.validation_check.libs.name_checks import ( check_if_kfp_pipeline_job_name_valid, ) @@ -191,6 +195,79 @@ logger = Logger() +# Map of start components to graph store compatibility checks to run +# Only run trainer checks when starting at or before Trainer +# Only run inferencer checks when starting at or before Inferencer +START_COMPONENT_TO_GRAPH_STORE_COMPATIBILITY_CHECKS = { + GiGLComponents.ConfigPopulator.value: [ + check_trainer_graph_store_compatibility, + check_inferencer_graph_store_compatibility, + ], + GiGLComponents.DataPreprocessor.value: [ + check_trainer_graph_store_compatibility, + check_inferencer_graph_store_compatibility, + ], + GiGLComponents.SubgraphSampler.value: [ + check_trainer_graph_store_compatibility, + check_inferencer_graph_store_compatibility, + ], + GiGLComponents.SplitGenerator.value: [ + check_trainer_graph_store_compatibility, + check_inferencer_graph_store_compatibility, + ], + GiGLComponents.Trainer.value: [ + check_trainer_graph_store_compatibility, + check_inferencer_graph_store_compatibility, + ], + GiGLComponents.Inferencer.value: [ + check_inferencer_graph_store_compatibility, + ], + # PostProcessor doesn't need graph store compatibility checks +} + +# Map of (start, stop) component tuples to graph store compatibility checks + +STOP_COMPONENT_TO_GRAPH_STORE_COMPATIBILITY_CHECKS_TO_SKIP = { + GiGLComponents.Trainer.value: [ + check_inferencer_graph_store_compatibility, + ], +} + + +def _run_gbml_and_resource_config_compatibility_checks( + start_at: str, + stop_after: Optional[str], + gbml_config_pb_wrapper: GbmlConfigPbWrapper, + resource_config_wrapper: GiglResourceConfigWrapper, +) -> None: + """ + Run compatibility checks between GbmlConfig and GiglResourceConfig. + + These checks verify that graph store mode configurations are consistent + across both the template config (GbmlConfig) and resource config (GiglResourceConfig). + + Args: + start_at: The component to start at. + stop_after: Optional component to stop after. + gbml_config_pb_wrapper: The GbmlConfig wrapper (template config). + resource_config_wrapper: The GiglResourceConfig wrapper (resource config). + """ + # Get the appropriate compatibility checks based on start/stop components + compatibility_checks = set( + START_COMPONENT_TO_GRAPH_STORE_COMPATIBILITY_CHECKS.get(start_at, []) + ) + if stop_after in STOP_COMPONENT_TO_GRAPH_STORE_COMPATIBILITY_CHECKS_TO_SKIP: + for skipped_check in STOP_COMPONENT_TO_GRAPH_STORE_COMPATIBILITY_CHECKS_TO_SKIP[ + stop_after + ]: + compatibility_checks.discard(skipped_check) + + for check in compatibility_checks: + check( + gbml_config_pb_wrapper=gbml_config_pb_wrapper, + resource_config_wrapper=resource_config_wrapper, + ) + def kfp_validation_checks( job_name: str, @@ -261,6 +338,15 @@ def kfp_validation_checks( f"Skipping resource config check {resource_config_check.__name__} because we are using live subgraph sampling backend." ) + # check compatibility between template config and resource config for graph store mode + # These checks ensure that if graph store mode is enabled in one config, it's also enabled in the other + _run_gbml_and_resource_config_compatibility_checks( + start_at=start_at, + stop_after=stop_after, + gbml_config_pb_wrapper=gbml_config_pb_wrapper, + resource_config_wrapper=resource_config_wrapper, + ) + # check if trained model file exist when skipping training if gbml_config_pb.shared_config.should_skip_training == True: assert_trained_model_exists(gbml_config_pb=gbml_config_pb) diff --git a/python/gigl/src/validation_check/libs/gbml_and_resource_config_compatibility_checks.py b/python/gigl/src/validation_check/libs/gbml_and_resource_config_compatibility_checks.py new file mode 100644 index 000000000..10f8bfc72 --- /dev/null +++ b/python/gigl/src/validation_check/libs/gbml_and_resource_config_compatibility_checks.py @@ -0,0 +1,167 @@ +""" +Compatibility checks between GbmlConfig (template config) and GiglResourceConfig (resource config). + +These checks ensure that graph store mode configurations are consistent across both configs. +If graph store mode is set up for trainer or inferencer in one config, it must be set up in the other. +""" + +from gigl.common.logger import Logger +from gigl.src.common.types.pb_wrappers.gbml_config import GbmlConfigPbWrapper +from gigl.src.common.types.pb_wrappers.gigl_resource_config import ( + GiglResourceConfigWrapper, +) +from snapchat.research.gbml import gigl_resource_config_pb2 + +logger = Logger() + + +def _gbml_config_has_trainer_graph_store( + gbml_config_pb_wrapper: GbmlConfigPbWrapper, +) -> bool: + """ + Check if the GbmlConfig has graph_store_storage_config set for trainer. + + Args: + gbml_config_pb_wrapper: The GbmlConfig wrapper to check. + + Returns: + True if graph_store_storage_config is set for trainer, False otherwise. + """ + trainer_config = gbml_config_pb_wrapper.gbml_config_pb.trainer_config + return trainer_config.HasField("graph_store_storage_config") + + +def _gbml_config_has_inferencer_graph_store( + gbml_config_pb_wrapper: GbmlConfigPbWrapper, +) -> bool: + """ + Check if the GbmlConfig has graph_store_storage_config set for inferencer. + + Args: + gbml_config_pb_wrapper: The GbmlConfig wrapper to check. + + Returns: + True if graph_store_storage_config is set for inferencer, False otherwise. + """ + inferencer_config = gbml_config_pb_wrapper.gbml_config_pb.inferencer_config + return inferencer_config.HasField("graph_store_storage_config") + + +def _resource_config_has_trainer_graph_store( + resource_config_wrapper: GiglResourceConfigWrapper, +) -> bool: + """ + Check if the GiglResourceConfig has VertexAiGraphStoreConfig set for trainer. + + Args: + resource_config_wrapper: The resource config wrapper to check. + + Returns: + True if VertexAiGraphStoreConfig is set for trainer, False otherwise. + """ + trainer_config = resource_config_wrapper.trainer_config + return isinstance(trainer_config, gigl_resource_config_pb2.VertexAiGraphStoreConfig) + + +def _resource_config_has_inferencer_graph_store( + resource_config_wrapper: GiglResourceConfigWrapper, +) -> bool: + """ + Check if the GiglResourceConfig has VertexAiGraphStoreConfig set for inferencer. + + Args: + resource_config_wrapper: The resource config wrapper to check. + + Returns: + True if VertexAiGraphStoreConfig is set for inferencer, False otherwise. + """ + inferencer_config = resource_config_wrapper.inferencer_config + return isinstance( + inferencer_config, gigl_resource_config_pb2.VertexAiGraphStoreConfig + ) + + +def check_trainer_graph_store_compatibility( + gbml_config_pb_wrapper: GbmlConfigPbWrapper, + resource_config_wrapper: GiglResourceConfigWrapper, +) -> None: + """ + Check that trainer graph store mode is consistently configured across both configs. + + If graph_store_storage_config is set in GbmlConfig.trainer_config, then + VertexAiGraphStoreConfig must be set in GiglResourceConfig.trainer_resource_config, + and vice versa. Also validates that storage_command is set when graph store mode is enabled. + + Args: + gbml_config_pb_wrapper: The GbmlConfig wrapper (template config). + resource_config_wrapper: The GiglResourceConfig wrapper (resource config). + + Raises: + AssertionError: If graph store configurations are not compatible or storage_command is missing. + """ + logger.info( + "Config validation check: trainer graph store compatibility between template and resource configs." + ) + + gbml_has_graph_store = _gbml_config_has_trainer_graph_store(gbml_config_pb_wrapper) + resource_has_graph_store = _resource_config_has_trainer_graph_store( + resource_config_wrapper + ) + + if gbml_has_graph_store and not resource_has_graph_store: + raise AssertionError( + "GbmlConfig.trainer_config.graph_store_storage_config is set, but " + "GiglResourceConfig.trainer_resource_config does not use VertexAiGraphStoreConfig. " + "Both configs must use graph store mode for trainer, or neither should." + ) + + if resource_has_graph_store and not gbml_has_graph_store: + raise AssertionError( + "GiglResourceConfig.trainer_resource_config uses VertexAiGraphStoreConfig, but " + "GbmlConfig.trainer_config.graph_store_storage_config is not set. " + "Both configs must use graph store mode for trainer, or neither should." + ) + + +def check_inferencer_graph_store_compatibility( + gbml_config_pb_wrapper: GbmlConfigPbWrapper, + resource_config_wrapper: GiglResourceConfigWrapper, +) -> None: + """ + Check that inferencer graph store mode is consistently configured across both configs. + + If graph_store_storage_config is set in GbmlConfig.inferencer_config, then + VertexAiGraphStoreConfig must be set in GiglResourceConfig.inferencer_resource_config, + and vice versa. Also validates that storage_command is set when graph store mode is enabled. + + Args: + gbml_config_pb_wrapper: The GbmlConfig wrapper (template config). + resource_config_wrapper: The GiglResourceConfig wrapper (resource config). + + Raises: + AssertionError: If graph store configurations are not compatible or storage_command is missing. + """ + logger.info( + "Config validation check: inferencer graph store compatibility between template and resource configs." + ) + + gbml_has_graph_store = _gbml_config_has_inferencer_graph_store( + gbml_config_pb_wrapper + ) + resource_has_graph_store = _resource_config_has_inferencer_graph_store( + resource_config_wrapper + ) + + if gbml_has_graph_store and not resource_has_graph_store: + raise AssertionError( + "GbmlConfig.inferencer_config.graph_store_storage_config is set, but " + "GiglResourceConfig.inferencer_resource_config does not use VertexAiGraphStoreConfig. " + "Both configs must use graph store mode for inferencer, or neither should." + ) + + if resource_has_graph_store and not gbml_has_graph_store: + raise AssertionError( + "GiglResourceConfig.inferencer_resource_config uses VertexAiGraphStoreConfig, but " + "GbmlConfig.inferencer_config.graph_store_storage_config is not set. " + "Both configs must use graph store mode for inferencer, or neither should." + ) diff --git a/python/gigl/src/validation_check/libs/resource_config_checks.py b/python/gigl/src/validation_check/libs/resource_config_checks.py index a60a66945..51cc60651 100644 --- a/python/gigl/src/validation_check/libs/resource_config_checks.py +++ b/python/gigl/src/validation_check/libs/resource_config_checks.py @@ -3,6 +3,7 @@ from google.cloud.aiplatform_v1.types.accelerator_type import AcceleratorType from gigl.common.logger import Logger +from gigl.src.common.types.pb_wrappers.gbml_config import GbmlConfigPbWrapper from gigl.src.common.types.pb_wrappers.gigl_resource_config import ( GiglResourceConfigWrapper, ) @@ -254,3 +255,53 @@ def _validate_machine_config( or {gigl_resource_config_pb2.VertexAiGraphStoreConfig.__name__}. Got {type(config)}""" ) + + +def check_if_trainer_graph_store_storage_command_valid( + gbml_config_pb_wrapper: GbmlConfigPbWrapper, +) -> None: + """ + Validates that storage_command is set when graph store mode is enabled for trainer. + + Args: + gbml_config_pb_wrapper: The GbmlConfig wrapper to check. + + Raises: + AssertionError: If graph store mode is enabled but storage_command is missing. + """ + logger.info( + "Config validation check: if trainer graph store storage_command is valid." + ) + trainer_config = gbml_config_pb_wrapper.gbml_config_pb.trainer_config + if trainer_config.HasField("graph_store_storage_config"): + storage_command = trainer_config.graph_store_storage_config.storage_command + if not storage_command: + raise AssertionError( + "GbmlConfig.trainer_config.graph_store_storage_config.storage_command must be set " + "when using graph store mode for trainer." + ) + + +def check_if_inferencer_graph_store_storage_command_valid( + gbml_config_pb_wrapper: GbmlConfigPbWrapper, +) -> None: + """ + Validates that storage_command is set when graph store mode is enabled for inferencer. + + Args: + gbml_config_pb_wrapper: The GbmlConfig wrapper to check. + + Raises: + AssertionError: If graph store mode is enabled but storage_command is missing. + """ + logger.info( + "Config validation check: if inferencer graph store storage_command is valid." + ) + inferencer_config = gbml_config_pb_wrapper.gbml_config_pb.inferencer_config + if inferencer_config.HasField("graph_store_storage_config"): + storage_command = inferencer_config.graph_store_storage_config.storage_command + if not storage_command: + raise AssertionError( + "GbmlConfig.inferencer_config.graph_store_storage_config.storage_command must be set " + "when using graph store mode for inferencer." + ) diff --git a/python/tests/unit/src/validation/lib/gbml_and_resource_config_compatibility_checks_test.py b/python/tests/unit/src/validation/lib/gbml_and_resource_config_compatibility_checks_test.py new file mode 100644 index 000000000..686fab9fb --- /dev/null +++ b/python/tests/unit/src/validation/lib/gbml_and_resource_config_compatibility_checks_test.py @@ -0,0 +1,403 @@ +import unittest + +from gigl.src.common.types.pb_wrappers.gbml_config import GbmlConfigPbWrapper +from gigl.src.common.types.pb_wrappers.gigl_resource_config import ( + GiglResourceConfigWrapper, +) +from gigl.src.validation_check.libs.gbml_and_resource_config_compatibility_checks import ( + check_inferencer_graph_store_compatibility, + check_trainer_graph_store_compatibility, +) +from snapchat.research.gbml import gbml_config_pb2, gigl_resource_config_pb2 + +# Helper functions for creating VertexAiGraphStoreConfig + + +def _create_vertex_ai_graph_store_config() -> ( + gigl_resource_config_pb2.VertexAiGraphStoreConfig +): + """Create a valid VertexAiGraphStoreConfig.""" + config = gigl_resource_config_pb2.VertexAiGraphStoreConfig() + # Graph store pool + config.graph_store_pool.machine_type = "n1-highmem-8" + config.graph_store_pool.gpu_type = "ACCELERATOR_TYPE_UNSPECIFIED" + config.graph_store_pool.gpu_limit = 0 + config.graph_store_pool.num_replicas = 2 + # Compute pool + config.compute_pool.machine_type = "n1-standard-16" + config.compute_pool.gpu_type = "NVIDIA_TESLA_T4" + config.compute_pool.gpu_limit = 2 + config.compute_pool.num_replicas = 3 + return config + + +def _create_vertex_ai_resource_config() -> ( + gigl_resource_config_pb2.VertexAiResourceConfig +): + """Create a valid VertexAiResourceConfig (non-graph store).""" + config = gigl_resource_config_pb2.VertexAiResourceConfig() + config.machine_type = "n1-standard-16" + config.gpu_type = "NVIDIA_TESLA_T4" + config.gpu_limit = 2 + config.num_replicas = 3 + return config + + +def _create_shared_resource_config( + config: gigl_resource_config_pb2.GiglResourceConfig, +) -> None: + """Populate shared resource config fields.""" + config.shared_resource_config.common_compute_config.project = "test-project" + config.shared_resource_config.common_compute_config.region = "us-central1" + config.shared_resource_config.common_compute_config.temp_assets_bucket = ( + "gs://test-temp" + ) + config.shared_resource_config.common_compute_config.temp_regional_assets_bucket = ( + "gs://test-temp-regional" + ) + config.shared_resource_config.common_compute_config.perm_assets_bucket = ( + "gs://test-perm" + ) + config.shared_resource_config.common_compute_config.temp_assets_bq_dataset_name = ( + "test_dataset" + ) + config.shared_resource_config.common_compute_config.embedding_bq_dataset_name = ( + "test_embeddings" + ) + config.shared_resource_config.common_compute_config.gcp_service_account_email = ( + "test@test-project.iam.gserviceaccount.com" + ) + config.shared_resource_config.common_compute_config.dataflow_runner = ( + "DataflowRunner" + ) + + +# Helper functions for creating GbmlConfig configurations + + +def _create_gbml_config_with_trainer_graph_store( + storage_command: str = "python -m gigl.distributed.graph_store.storage_main", +) -> GbmlConfigPbWrapper: + """Create a GbmlConfig with graph_store_storage_config set for trainer.""" + gbml_config = gbml_config_pb2.GbmlConfig() + gbml_config.trainer_config.graph_store_storage_config.storage_command = ( + storage_command + ) + return GbmlConfigPbWrapper(gbml_config_pb=gbml_config) + + +def _create_gbml_config_without_trainer_graph_store() -> GbmlConfigPbWrapper: + """Create a GbmlConfig without graph_store_storage_config for trainer.""" + gbml_config = gbml_config_pb2.GbmlConfig() + gbml_config.trainer_config.trainer_args["some_arg"] = "some_value" + return GbmlConfigPbWrapper(gbml_config_pb=gbml_config) + + +def _create_gbml_config_with_inferencer_graph_store( + storage_command: str = "python -m gigl.distributed.graph_store.storage_main", +) -> GbmlConfigPbWrapper: + """Create a GbmlConfig with graph_store_storage_config set for inferencer.""" + gbml_config = gbml_config_pb2.GbmlConfig() + gbml_config.inferencer_config.graph_store_storage_config.storage_command = ( + storage_command + ) + return GbmlConfigPbWrapper(gbml_config_pb=gbml_config) + + +def _create_gbml_config_without_inferencer_graph_store() -> GbmlConfigPbWrapper: + """Create a GbmlConfig without graph_store_storage_config for inferencer.""" + gbml_config = gbml_config_pb2.GbmlConfig() + gbml_config.inferencer_config.inferencer_args["some_arg"] = "some_value" + return GbmlConfigPbWrapper(gbml_config_pb=gbml_config) + + +def _create_gbml_config_with_both_graph_stores( + storage_command: str = "python -m gigl.distributed.graph_store.storage_main", +) -> GbmlConfigPbWrapper: + """Create a GbmlConfig with graph_store_storage_config for both trainer and inferencer.""" + gbml_config = gbml_config_pb2.GbmlConfig() + gbml_config.trainer_config.graph_store_storage_config.storage_command = ( + storage_command + ) + gbml_config.inferencer_config.graph_store_storage_config.storage_command = ( + storage_command + ) + return GbmlConfigPbWrapper(gbml_config_pb=gbml_config) + + +def _create_gbml_config_without_graph_stores() -> GbmlConfigPbWrapper: + """Create a GbmlConfig without graph_store_storage_config for both trainer and inferencer.""" + gbml_config = gbml_config_pb2.GbmlConfig() + gbml_config.trainer_config.trainer_args["some_arg"] = "some_value" + gbml_config.inferencer_config.inferencer_args["some_arg"] = "some_value" + return GbmlConfigPbWrapper(gbml_config_pb=gbml_config) + + +# Helper functions for creating GiglResourceConfig configurations + + +def _create_resource_config_with_trainer_graph_store() -> GiglResourceConfigWrapper: + """Create a GiglResourceConfig with VertexAiGraphStoreConfig for trainer.""" + config = gigl_resource_config_pb2.GiglResourceConfig() + _create_shared_resource_config(config) + + # Trainer with VertexAiGraphStoreConfig + config.trainer_resource_config.vertex_ai_graph_store_trainer_config.CopyFrom( + _create_vertex_ai_graph_store_config() + ) + # Inferencer with standard config + config.inferencer_resource_config.vertex_ai_inferencer_config.CopyFrom( + _create_vertex_ai_resource_config() + ) + return GiglResourceConfigWrapper(resource_config=config) + + +def _create_resource_config_without_trainer_graph_store() -> GiglResourceConfigWrapper: + """Create a GiglResourceConfig without VertexAiGraphStoreConfig for trainer.""" + config = gigl_resource_config_pb2.GiglResourceConfig() + _create_shared_resource_config(config) + + # Trainer with standard config + config.trainer_resource_config.vertex_ai_trainer_config.CopyFrom( + _create_vertex_ai_resource_config() + ) + # Inferencer with standard config + config.inferencer_resource_config.vertex_ai_inferencer_config.CopyFrom( + _create_vertex_ai_resource_config() + ) + return GiglResourceConfigWrapper(resource_config=config) + + +def _create_resource_config_with_inferencer_graph_store() -> GiglResourceConfigWrapper: + """Create a GiglResourceConfig with VertexAiGraphStoreConfig for inferencer.""" + config = gigl_resource_config_pb2.GiglResourceConfig() + _create_shared_resource_config(config) + + # Trainer with standard config + config.trainer_resource_config.vertex_ai_trainer_config.CopyFrom( + _create_vertex_ai_resource_config() + ) + # Inferencer with VertexAiGraphStoreConfig + config.inferencer_resource_config.vertex_ai_graph_store_inferencer_config.CopyFrom( + _create_vertex_ai_graph_store_config() + ) + return GiglResourceConfigWrapper(resource_config=config) + + +def _create_resource_config_without_inferencer_graph_store() -> ( + GiglResourceConfigWrapper +): + """Create a GiglResourceConfig without VertexAiGraphStoreConfig for inferencer.""" + return _create_resource_config_without_trainer_graph_store() + + +def _create_resource_config_with_both_graph_stores() -> GiglResourceConfigWrapper: + """Create a GiglResourceConfig with VertexAiGraphStoreConfig for both trainer and inferencer.""" + config = gigl_resource_config_pb2.GiglResourceConfig() + _create_shared_resource_config(config) + + # Trainer with VertexAiGraphStoreConfig + config.trainer_resource_config.vertex_ai_graph_store_trainer_config.CopyFrom( + _create_vertex_ai_graph_store_config() + ) + # Inferencer with VertexAiGraphStoreConfig + config.inferencer_resource_config.vertex_ai_graph_store_inferencer_config.CopyFrom( + _create_vertex_ai_graph_store_config() + ) + return GiglResourceConfigWrapper(resource_config=config) + + +def _create_resource_config_without_graph_stores() -> GiglResourceConfigWrapper: + """Create a GiglResourceConfig without VertexAiGraphStoreConfig for both trainer and inferencer.""" + return _create_resource_config_without_trainer_graph_store() + + +# Test Classes + + +class TestTrainerGraphStoreCompatibility(unittest.TestCase): + """Test suite for trainer graph store compatibility checks.""" + + def test_both_have_trainer_graph_store(self): + """Test that both configs having trainer graph store passes validation.""" + gbml_config = _create_gbml_config_with_trainer_graph_store() + resource_config = _create_resource_config_with_trainer_graph_store() + # Should not raise any exception + check_trainer_graph_store_compatibility( + gbml_config_pb_wrapper=gbml_config, + resource_config_wrapper=resource_config, + ) + + def test_neither_has_trainer_graph_store(self): + """Test that neither config having trainer graph store passes validation.""" + gbml_config = _create_gbml_config_without_trainer_graph_store() + resource_config = _create_resource_config_without_trainer_graph_store() + # Should not raise any exception + check_trainer_graph_store_compatibility( + gbml_config_pb_wrapper=gbml_config, + resource_config_wrapper=resource_config, + ) + + def test_template_has_trainer_graph_store_resource_does_not(self): + """Test that template having graph store but resource not raises an assertion error.""" + gbml_config = _create_gbml_config_with_trainer_graph_store() + resource_config = _create_resource_config_without_trainer_graph_store() + with self.assertRaises(AssertionError): + check_trainer_graph_store_compatibility( + gbml_config_pb_wrapper=gbml_config, + resource_config_wrapper=resource_config, + ) + + def test_resource_has_trainer_graph_store_template_does_not(self): + """Test that resource having graph store but template not raises an assertion error.""" + gbml_config = _create_gbml_config_without_trainer_graph_store() + resource_config = _create_resource_config_with_trainer_graph_store() + with self.assertRaises(AssertionError): + check_trainer_graph_store_compatibility( + gbml_config_pb_wrapper=gbml_config, + resource_config_wrapper=resource_config, + ) + + +class TestInferencerGraphStoreCompatibility(unittest.TestCase): + """Test suite for inferencer graph store compatibility checks.""" + + def test_both_have_inferencer_graph_store(self): + """Test that both configs having inferencer graph store passes validation.""" + gbml_config = _create_gbml_config_with_inferencer_graph_store() + resource_config = _create_resource_config_with_inferencer_graph_store() + # Should not raise any exception + check_inferencer_graph_store_compatibility( + gbml_config_pb_wrapper=gbml_config, + resource_config_wrapper=resource_config, + ) + + def test_neither_has_inferencer_graph_store(self): + """Test that neither config having inferencer graph store passes validation.""" + gbml_config = _create_gbml_config_without_inferencer_graph_store() + resource_config = _create_resource_config_without_inferencer_graph_store() + # Should not raise any exception + check_inferencer_graph_store_compatibility( + gbml_config_pb_wrapper=gbml_config, + resource_config_wrapper=resource_config, + ) + + def test_template_has_inferencer_graph_store_resource_does_not(self): + """Test that template having graph store but resource not raises an assertion error.""" + gbml_config = _create_gbml_config_with_inferencer_graph_store() + resource_config = _create_resource_config_without_inferencer_graph_store() + with self.assertRaises(AssertionError): + check_inferencer_graph_store_compatibility( + gbml_config_pb_wrapper=gbml_config, + resource_config_wrapper=resource_config, + ) + + def test_resource_has_inferencer_graph_store_template_does_not(self): + """Test that resource having graph store but template not raises an assertion error.""" + gbml_config = _create_gbml_config_without_inferencer_graph_store() + resource_config = _create_resource_config_with_inferencer_graph_store() + with self.assertRaises(AssertionError): + check_inferencer_graph_store_compatibility( + gbml_config_pb_wrapper=gbml_config, + resource_config_wrapper=resource_config, + ) + + +class TestMixedGraphStoreConfigurations(unittest.TestCase): + """Test suite for mixed graph store configuration scenarios.""" + + def test_both_have_all_graph_stores(self): + """Test that both configs having all graph stores passes validation.""" + gbml_config = _create_gbml_config_with_both_graph_stores() + resource_config = _create_resource_config_with_both_graph_stores() + # Should not raise any exception for trainer + check_trainer_graph_store_compatibility( + gbml_config_pb_wrapper=gbml_config, + resource_config_wrapper=resource_config, + ) + # Should not raise any exception for inferencer + check_inferencer_graph_store_compatibility( + gbml_config_pb_wrapper=gbml_config, + resource_config_wrapper=resource_config, + ) + + def test_neither_has_any_graph_stores(self): + """Test that neither config having any graph stores passes validation.""" + gbml_config = _create_gbml_config_without_graph_stores() + resource_config = _create_resource_config_without_graph_stores() + # Should not raise any exception for trainer + check_trainer_graph_store_compatibility( + gbml_config_pb_wrapper=gbml_config, + resource_config_wrapper=resource_config, + ) + # Should not raise any exception for inferencer + check_inferencer_graph_store_compatibility( + gbml_config_pb_wrapper=gbml_config, + resource_config_wrapper=resource_config, + ) + + def test_trainer_graph_store_only_compatible(self): + """Test trainer graph store only configuration is compatible.""" + gbml_config = _create_gbml_config_with_trainer_graph_store() + resource_config = _create_resource_config_with_trainer_graph_store() + # Should not raise any exception for trainer + check_trainer_graph_store_compatibility( + gbml_config_pb_wrapper=gbml_config, + resource_config_wrapper=resource_config, + ) + # Should not raise any exception for inferencer (neither has it) + check_inferencer_graph_store_compatibility( + gbml_config_pb_wrapper=gbml_config, + resource_config_wrapper=resource_config, + ) + + def test_inferencer_graph_store_only_compatible(self): + """Test inferencer graph store only configuration is compatible.""" + gbml_config = _create_gbml_config_with_inferencer_graph_store() + resource_config = _create_resource_config_with_inferencer_graph_store() + # Should not raise any exception for trainer (neither has it) + check_trainer_graph_store_compatibility( + gbml_config_pb_wrapper=gbml_config, + resource_config_wrapper=resource_config, + ) + # Should not raise any exception for inferencer + check_inferencer_graph_store_compatibility( + gbml_config_pb_wrapper=gbml_config, + resource_config_wrapper=resource_config, + ) + + def test_template_has_both_resource_has_trainer_only(self): + """Test that template having both but resource having only trainer raises an error for inferencer.""" + gbml_config = _create_gbml_config_with_both_graph_stores() + resource_config = _create_resource_config_with_trainer_graph_store() + # Should not raise any exception for trainer + check_trainer_graph_store_compatibility( + gbml_config_pb_wrapper=gbml_config, + resource_config_wrapper=resource_config, + ) + # Should raise an assertion error for inferencer + with self.assertRaises(AssertionError): + check_inferencer_graph_store_compatibility( + gbml_config_pb_wrapper=gbml_config, + resource_config_wrapper=resource_config, + ) + + def test_template_has_both_resource_has_inferencer_only(self): + """Test that template having both but resource having only inferencer raises an error for trainer.""" + gbml_config = _create_gbml_config_with_both_graph_stores() + resource_config = _create_resource_config_with_inferencer_graph_store() + # Should raise an assertion error for trainer + with self.assertRaises(AssertionError): + check_trainer_graph_store_compatibility( + gbml_config_pb_wrapper=gbml_config, + resource_config_wrapper=resource_config, + ) + # Should not raise any exception for inferencer + check_inferencer_graph_store_compatibility( + gbml_config_pb_wrapper=gbml_config, + resource_config_wrapper=resource_config, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/python/tests/unit/src/validation/lib/resource_config_checks_test.py b/python/tests/unit/src/validation/lib/resource_config_checks_test.py index 20b941f72..3bddf1269 100644 --- a/python/tests/unit/src/validation/lib/resource_config_checks_test.py +++ b/python/tests/unit/src/validation/lib/resource_config_checks_test.py @@ -1,18 +1,21 @@ import unittest +from gigl.src.common.types.pb_wrappers.gbml_config import GbmlConfigPbWrapper from gigl.src.validation_check.libs.resource_config_checks import ( _check_if_dataflow_resource_config_valid, _check_if_spark_resource_config_valid, _validate_accelerator_type, _validate_machine_config, + check_if_inferencer_graph_store_storage_command_valid, check_if_inferencer_resource_config_valid, check_if_preprocessor_resource_config_valid, check_if_shared_resource_config_valid, check_if_split_generator_resource_config_valid, check_if_subgraph_sampler_resource_config_valid, + check_if_trainer_graph_store_storage_command_valid, check_if_trainer_resource_config_valid, ) -from snapchat.research.gbml import gigl_resource_config_pb2 +from snapchat.research.gbml import gbml_config_pb2, gigl_resource_config_pb2 # Helper functions for creating valid configurations @@ -750,5 +753,90 @@ def test_valid_vertex_ai_graph_store_config(self): _validate_machine_config(config) +# Helper functions for creating GbmlConfig configurations + + +def _create_gbml_config_with_trainer_graph_store( + storage_command: str = "python -m gigl.distributed.graph_store.storage_main", +) -> GbmlConfigPbWrapper: + """Create a GbmlConfig with graph_store_storage_config set for trainer.""" + gbml_config = gbml_config_pb2.GbmlConfig() + gbml_config.trainer_config.graph_store_storage_config.storage_command = ( + storage_command + ) + return GbmlConfigPbWrapper(gbml_config_pb=gbml_config) + + +def _create_gbml_config_without_trainer_graph_store() -> GbmlConfigPbWrapper: + """Create a GbmlConfig without graph_store_storage_config for trainer.""" + gbml_config = gbml_config_pb2.GbmlConfig() + gbml_config.trainer_config.trainer_args["some_arg"] = "some_value" + return GbmlConfigPbWrapper(gbml_config_pb=gbml_config) + + +def _create_gbml_config_with_inferencer_graph_store( + storage_command: str = "python -m gigl.distributed.graph_store.storage_main", +) -> GbmlConfigPbWrapper: + """Create a GbmlConfig with graph_store_storage_config set for inferencer.""" + gbml_config = gbml_config_pb2.GbmlConfig() + gbml_config.inferencer_config.graph_store_storage_config.storage_command = ( + storage_command + ) + return GbmlConfigPbWrapper(gbml_config_pb=gbml_config) + + +def _create_gbml_config_without_inferencer_graph_store() -> GbmlConfigPbWrapper: + """Create a GbmlConfig without graph_store_storage_config for inferencer.""" + gbml_config = gbml_config_pb2.GbmlConfig() + gbml_config.inferencer_config.inferencer_args["some_arg"] = "some_value" + return GbmlConfigPbWrapper(gbml_config_pb=gbml_config) + + +class TestTrainerGraphStoreStorageCommand(unittest.TestCase): + """Test suite for trainer graph store storage_command validation.""" + + def test_valid_storage_command(self): + """Test that a valid storage_command passes validation.""" + gbml_config = _create_gbml_config_with_trainer_graph_store() + # Should not raise any exception + check_if_trainer_graph_store_storage_command_valid(gbml_config) + + def test_missing_storage_command(self): + """Test that missing storage_command raises an assertion error.""" + gbml_config = _create_gbml_config_with_trainer_graph_store(storage_command="") + with self.assertRaises(AssertionError): + check_if_trainer_graph_store_storage_command_valid(gbml_config) + + def test_no_graph_store_config(self): + """Test that no graph store config passes validation (nothing to check).""" + gbml_config = _create_gbml_config_without_trainer_graph_store() + # Should not raise any exception - no graph store means nothing to validate + check_if_trainer_graph_store_storage_command_valid(gbml_config) + + +class TestInferencerGraphStoreStorageCommand(unittest.TestCase): + """Test suite for inferencer graph store storage_command validation.""" + + def test_valid_storage_command(self): + """Test that a valid storage_command passes validation.""" + gbml_config = _create_gbml_config_with_inferencer_graph_store() + # Should not raise any exception + check_if_inferencer_graph_store_storage_command_valid(gbml_config) + + def test_missing_storage_command(self): + """Test that missing storage_command raises an assertion error.""" + gbml_config = _create_gbml_config_with_inferencer_graph_store( + storage_command="" + ) + with self.assertRaises(AssertionError): + check_if_inferencer_graph_store_storage_command_valid(gbml_config) + + def test_no_graph_store_config(self): + """Test that no graph store config passes validation (nothing to check).""" + gbml_config = _create_gbml_config_without_inferencer_graph_store() + # Should not raise any exception - no graph store means nothing to validate + check_if_inferencer_graph_store_storage_command_valid(gbml_config) + + if __name__ == "__main__": unittest.main() From 7aa4b79366e3ad05eb80c61598dc3df61b21dc65 Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Fri, 23 Jan 2026 00:34:33 +0000 Subject: [PATCH 2/6] eveolve --- .../libs/resource_config_checks.py | 4 ++-- ..._resource_config_compatibility_checks_test.py | 16 ++++------------ .../lib/resource_config_checks_test.py | 8 ++------ 3 files changed, 8 insertions(+), 20 deletions(-) diff --git a/python/gigl/src/validation_check/libs/resource_config_checks.py b/python/gigl/src/validation_check/libs/resource_config_checks.py index 51cc60651..8e1137726 100644 --- a/python/gigl/src/validation_check/libs/resource_config_checks.py +++ b/python/gigl/src/validation_check/libs/resource_config_checks.py @@ -274,7 +274,7 @@ def check_if_trainer_graph_store_storage_command_valid( ) trainer_config = gbml_config_pb_wrapper.gbml_config_pb.trainer_config if trainer_config.HasField("graph_store_storage_config"): - storage_command = trainer_config.graph_store_storage_config.storage_command + storage_command = trainer_config.graph_store_storage_config.command if not storage_command: raise AssertionError( "GbmlConfig.trainer_config.graph_store_storage_config.storage_command must be set " @@ -299,7 +299,7 @@ def check_if_inferencer_graph_store_storage_command_valid( ) inferencer_config = gbml_config_pb_wrapper.gbml_config_pb.inferencer_config if inferencer_config.HasField("graph_store_storage_config"): - storage_command = inferencer_config.graph_store_storage_config.storage_command + storage_command = inferencer_config.graph_store_storage_config.command if not storage_command: raise AssertionError( "GbmlConfig.inferencer_config.graph_store_storage_config.storage_command must be set " diff --git a/python/tests/unit/src/validation/lib/gbml_and_resource_config_compatibility_checks_test.py b/python/tests/unit/src/validation/lib/gbml_and_resource_config_compatibility_checks_test.py index 686fab9fb..07b0e6d46 100644 --- a/python/tests/unit/src/validation/lib/gbml_and_resource_config_compatibility_checks_test.py +++ b/python/tests/unit/src/validation/lib/gbml_and_resource_config_compatibility_checks_test.py @@ -80,9 +80,7 @@ def _create_gbml_config_with_trainer_graph_store( ) -> GbmlConfigPbWrapper: """Create a GbmlConfig with graph_store_storage_config set for trainer.""" gbml_config = gbml_config_pb2.GbmlConfig() - gbml_config.trainer_config.graph_store_storage_config.storage_command = ( - storage_command - ) + gbml_config.trainer_config.graph_store_storage_config.command = storage_command return GbmlConfigPbWrapper(gbml_config_pb=gbml_config) @@ -98,9 +96,7 @@ def _create_gbml_config_with_inferencer_graph_store( ) -> GbmlConfigPbWrapper: """Create a GbmlConfig with graph_store_storage_config set for inferencer.""" gbml_config = gbml_config_pb2.GbmlConfig() - gbml_config.inferencer_config.graph_store_storage_config.storage_command = ( - storage_command - ) + gbml_config.inferencer_config.graph_store_storage_config.command = storage_command return GbmlConfigPbWrapper(gbml_config_pb=gbml_config) @@ -116,12 +112,8 @@ def _create_gbml_config_with_both_graph_stores( ) -> GbmlConfigPbWrapper: """Create a GbmlConfig with graph_store_storage_config for both trainer and inferencer.""" gbml_config = gbml_config_pb2.GbmlConfig() - gbml_config.trainer_config.graph_store_storage_config.storage_command = ( - storage_command - ) - gbml_config.inferencer_config.graph_store_storage_config.storage_command = ( - storage_command - ) + gbml_config.trainer_config.graph_store_storage_config.command = storage_command + gbml_config.inferencer_config.graph_store_storage_config.command = storage_command return GbmlConfigPbWrapper(gbml_config_pb=gbml_config) diff --git a/python/tests/unit/src/validation/lib/resource_config_checks_test.py b/python/tests/unit/src/validation/lib/resource_config_checks_test.py index 3bddf1269..eff49aaab 100644 --- a/python/tests/unit/src/validation/lib/resource_config_checks_test.py +++ b/python/tests/unit/src/validation/lib/resource_config_checks_test.py @@ -761,9 +761,7 @@ def _create_gbml_config_with_trainer_graph_store( ) -> GbmlConfigPbWrapper: """Create a GbmlConfig with graph_store_storage_config set for trainer.""" gbml_config = gbml_config_pb2.GbmlConfig() - gbml_config.trainer_config.graph_store_storage_config.storage_command = ( - storage_command - ) + gbml_config.trainer_config.graph_store_storage_config.command = storage_command return GbmlConfigPbWrapper(gbml_config_pb=gbml_config) @@ -779,9 +777,7 @@ def _create_gbml_config_with_inferencer_graph_store( ) -> GbmlConfigPbWrapper: """Create a GbmlConfig with graph_store_storage_config set for inferencer.""" gbml_config = gbml_config_pb2.GbmlConfig() - gbml_config.inferencer_config.graph_store_storage_config.storage_command = ( - storage_command - ) + gbml_config.inferencer_config.graph_store_storage_config.command = storage_command return GbmlConfigPbWrapper(gbml_config_pb=gbml_config) From ec80ce0af6abffdfd67dc5ecea235749c987b6e5 Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Fri, 30 Jan 2026 17:04:43 +0000 Subject: [PATCH 3/6] update --- ...nd_resource_config_compatibility_checks.py | 108 +++++++----------- 1 file changed, 39 insertions(+), 69 deletions(-) diff --git a/gigl/src/validation_check/libs/gbml_and_resource_config_compatibility_checks.py b/gigl/src/validation_check/libs/gbml_and_resource_config_compatibility_checks.py index 10f8bfc72..796e62ec1 100644 --- a/gigl/src/validation_check/libs/gbml_and_resource_config_compatibility_checks.py +++ b/gigl/src/validation_check/libs/gbml_and_resource_config_compatibility_checks.py @@ -5,6 +5,10 @@ If graph store mode is set up for trainer or inferencer in one config, it must be set up in the other. """ +from typing import Literal + +from google.protobuf.message import Message + from gigl.common.logger import Logger from gigl.src.common.types.pb_wrappers.gbml_config import GbmlConfigPbWrapper from gigl.src.common.types.pb_wrappers.gigl_resource_config import ( @@ -15,24 +19,9 @@ logger = Logger() -def _gbml_config_has_trainer_graph_store( - gbml_config_pb_wrapper: GbmlConfigPbWrapper, -) -> bool: - """ - Check if the GbmlConfig has graph_store_storage_config set for trainer. - - Args: - gbml_config_pb_wrapper: The GbmlConfig wrapper to check. - - Returns: - True if graph_store_storage_config is set for trainer, False otherwise. - """ - trainer_config = gbml_config_pb_wrapper.gbml_config_pb.trainer_config - return trainer_config.HasField("graph_store_storage_config") - - -def _gbml_config_has_inferencer_graph_store( +def _gbml_config_has_graph_store( gbml_config_pb_wrapper: GbmlConfigPbWrapper, + component: Literal["trainer", "inferencer"], ) -> bool: """ Check if the GbmlConfig has graph_store_storage_config set for inferencer. @@ -43,15 +32,23 @@ def _gbml_config_has_inferencer_graph_store( Returns: True if graph_store_storage_config is set for inferencer, False otherwise. """ - inferencer_config = gbml_config_pb_wrapper.gbml_config_pb.inferencer_config - return inferencer_config.HasField("graph_store_storage_config") + if component == "inferencer": + config: Message = gbml_config_pb_wrapper.gbml_config_pb.inferencer_config + elif component == "trainer": + config = gbml_config_pb_wrapper.gbml_config_pb.trainer_config + else: + raise ValueError( + f"Invalid component: {component}. Must be 'inferencer' or 'trainer'." + ) + return config.HasField("graph_store_storage_config") -def _resource_config_has_trainer_graph_store( +def _resource_config_has_graph_store( resource_config_wrapper: GiglResourceConfigWrapper, + component: Literal["trainer", "inferencer"], ) -> bool: """ - Check if the GiglResourceConfig has VertexAiGraphStoreConfig set for trainer. + Check if the GiglResourceConfig has VertexAiGraphStoreConfig set for the given component. Args: resource_config_wrapper: The resource config wrapper to check. @@ -59,26 +56,15 @@ def _resource_config_has_trainer_graph_store( Returns: True if VertexAiGraphStoreConfig is set for trainer, False otherwise. """ - trainer_config = resource_config_wrapper.trainer_config - return isinstance(trainer_config, gigl_resource_config_pb2.VertexAiGraphStoreConfig) - - -def _resource_config_has_inferencer_graph_store( - resource_config_wrapper: GiglResourceConfigWrapper, -) -> bool: - """ - Check if the GiglResourceConfig has VertexAiGraphStoreConfig set for inferencer. - - Args: - resource_config_wrapper: The resource config wrapper to check. - - Returns: - True if VertexAiGraphStoreConfig is set for inferencer, False otherwise. - """ - inferencer_config = resource_config_wrapper.inferencer_config - return isinstance( - inferencer_config, gigl_resource_config_pb2.VertexAiGraphStoreConfig - ) + if component == "trainer": + config: Message = resource_config_wrapper.trainer_config + elif component == "inferencer": + config = resource_config_wrapper.inferencer_config + else: + raise ValueError( + f"Invalid component: {component}. Must be 'trainer' or 'inferencer'." + ) + return isinstance(config, gigl_resource_config_pb2.VertexAiGraphStoreConfig) def check_trainer_graph_store_compatibility( @@ -103,23 +89,16 @@ def check_trainer_graph_store_compatibility( "Config validation check: trainer graph store compatibility between template and resource configs." ) - gbml_has_graph_store = _gbml_config_has_trainer_graph_store(gbml_config_pb_wrapper) - resource_has_graph_store = _resource_config_has_trainer_graph_store( - resource_config_wrapper + gbml_has_graph_store = _gbml_config_has_graph_store( + gbml_config_pb_wrapper, "trainer" + ) + resource_has_graph_store = _resource_config_has_graph_store( + resource_config_wrapper, "trainer" ) - if gbml_has_graph_store and not resource_has_graph_store: - raise AssertionError( - "GbmlConfig.trainer_config.graph_store_storage_config is set, but " - "GiglResourceConfig.trainer_resource_config does not use VertexAiGraphStoreConfig. " - "Both configs must use graph store mode for trainer, or neither should." - ) - - if resource_has_graph_store and not gbml_has_graph_store: + if gbml_has_graph_store ^ resource_has_graph_store: raise AssertionError( - "GiglResourceConfig.trainer_resource_config uses VertexAiGraphStoreConfig, but " - "GbmlConfig.trainer_config.graph_store_storage_config is not set. " - "Both configs must use graph store mode for trainer, or neither should." + f"If one of GbmlConfig.trainer_config.graph_store_storage_config or GiglResourceConfig.trainer_resource_config is set, the other must also be set. GbmlConfig.trainer_config.graph_store_storage_config is set: {gbml_has_graph_store}, GiglResourceConfig.trainer_resource_config is set: {resource_has_graph_store}." ) @@ -145,23 +124,14 @@ def check_inferencer_graph_store_compatibility( "Config validation check: inferencer graph store compatibility between template and resource configs." ) - gbml_has_graph_store = _gbml_config_has_inferencer_graph_store( - gbml_config_pb_wrapper + gbml_has_graph_store = _gbml_config_has_graph_store( + gbml_config_pb_wrapper, "inferencer" ) - resource_has_graph_store = _resource_config_has_inferencer_graph_store( - resource_config_wrapper + resource_has_graph_store = _resource_config_has_graph_store( + resource_config_wrapper, "inferencer" ) if gbml_has_graph_store and not resource_has_graph_store: raise AssertionError( - "GbmlConfig.inferencer_config.graph_store_storage_config is set, but " - "GiglResourceConfig.inferencer_resource_config does not use VertexAiGraphStoreConfig. " - "Both configs must use graph store mode for inferencer, or neither should." - ) - - if resource_has_graph_store and not gbml_has_graph_store: - raise AssertionError( - "GiglResourceConfig.inferencer_resource_config uses VertexAiGraphStoreConfig, but " - "GbmlConfig.inferencer_config.graph_store_storage_config is not set. " - "Both configs must use graph store mode for inferencer, or neither should." + f"If one of GbmlConfig.inferencer_config.graph_store_storage_config or GiglResourceConfig.inferencer_resource_config is set, the other must also be set. GbmlConfig.inferencer_config.graph_store_storage_config is set: {gbml_has_graph_store}, GiglResourceConfig.inferencer_resource_config is set: {resource_has_graph_store}." ) From e9c5f9a5ecd590a88433bf01d2e9ccffcc31a25e Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Fri, 30 Jan 2026 17:13:25 +0000 Subject: [PATCH 4/6] fixes --- ...nd_resource_config_compatibility_checks.py | 2 +- ...source_config_compatibility_checks_test.py | 226 ++---------------- 2 files changed, 25 insertions(+), 203 deletions(-) diff --git a/gigl/src/validation_check/libs/gbml_and_resource_config_compatibility_checks.py b/gigl/src/validation_check/libs/gbml_and_resource_config_compatibility_checks.py index 796e62ec1..84dd75831 100644 --- a/gigl/src/validation_check/libs/gbml_and_resource_config_compatibility_checks.py +++ b/gigl/src/validation_check/libs/gbml_and_resource_config_compatibility_checks.py @@ -131,7 +131,7 @@ def check_inferencer_graph_store_compatibility( resource_config_wrapper, "inferencer" ) - if gbml_has_graph_store and not resource_has_graph_store: + if gbml_has_graph_store ^ resource_has_graph_store: raise AssertionError( f"If one of GbmlConfig.inferencer_config.graph_store_storage_config or GiglResourceConfig.inferencer_resource_config is set, the other must also be set. GbmlConfig.inferencer_config.graph_store_storage_config is set: {gbml_has_graph_store}, GiglResourceConfig.inferencer_resource_config is set: {resource_has_graph_store}." ) diff --git a/tests/unit/src/validation/lib/gbml_and_resource_config_compatibility_checks_test.py b/tests/unit/src/validation/lib/gbml_and_resource_config_compatibility_checks_test.py index 07b0e6d46..49966df63 100644 --- a/tests/unit/src/validation/lib/gbml_and_resource_config_compatibility_checks_test.py +++ b/tests/unit/src/validation/lib/gbml_and_resource_config_compatibility_checks_test.py @@ -75,42 +75,10 @@ def _create_shared_resource_config( # Helper functions for creating GbmlConfig configurations -def _create_gbml_config_with_trainer_graph_store( - storage_command: str = "python -m gigl.distributed.graph_store.storage_main", -) -> GbmlConfigPbWrapper: - """Create a GbmlConfig with graph_store_storage_config set for trainer.""" - gbml_config = gbml_config_pb2.GbmlConfig() - gbml_config.trainer_config.graph_store_storage_config.command = storage_command - return GbmlConfigPbWrapper(gbml_config_pb=gbml_config) - - -def _create_gbml_config_without_trainer_graph_store() -> GbmlConfigPbWrapper: - """Create a GbmlConfig without graph_store_storage_config for trainer.""" - gbml_config = gbml_config_pb2.GbmlConfig() - gbml_config.trainer_config.trainer_args["some_arg"] = "some_value" - return GbmlConfigPbWrapper(gbml_config_pb=gbml_config) - - -def _create_gbml_config_with_inferencer_graph_store( - storage_command: str = "python -m gigl.distributed.graph_store.storage_main", -) -> GbmlConfigPbWrapper: - """Create a GbmlConfig with graph_store_storage_config set for inferencer.""" - gbml_config = gbml_config_pb2.GbmlConfig() - gbml_config.inferencer_config.graph_store_storage_config.command = storage_command - return GbmlConfigPbWrapper(gbml_config_pb=gbml_config) - - -def _create_gbml_config_without_inferencer_graph_store() -> GbmlConfigPbWrapper: - """Create a GbmlConfig without graph_store_storage_config for inferencer.""" - gbml_config = gbml_config_pb2.GbmlConfig() - gbml_config.inferencer_config.inferencer_args["some_arg"] = "some_value" - return GbmlConfigPbWrapper(gbml_config_pb=gbml_config) - - def _create_gbml_config_with_both_graph_stores( storage_command: str = "python -m gigl.distributed.graph_store.storage_main", ) -> GbmlConfigPbWrapper: - """Create a GbmlConfig with graph_store_storage_config for both trainer and inferencer.""" + """Create a GbmlConfig with graph_store_storage_config set for both trainer and inferencer.""" gbml_config = gbml_config_pb2.GbmlConfig() gbml_config.trainer_config.graph_store_storage_config.command = storage_command gbml_config.inferencer_config.graph_store_storage_config.command = storage_command @@ -118,18 +86,15 @@ def _create_gbml_config_with_both_graph_stores( def _create_gbml_config_without_graph_stores() -> GbmlConfigPbWrapper: - """Create a GbmlConfig without graph_store_storage_config for both trainer and inferencer.""" + """Create a GbmlConfig without graph_store_storage_config for trainer or inferencer.""" gbml_config = gbml_config_pb2.GbmlConfig() gbml_config.trainer_config.trainer_args["some_arg"] = "some_value" gbml_config.inferencer_config.inferencer_args["some_arg"] = "some_value" return GbmlConfigPbWrapper(gbml_config_pb=gbml_config) -# Helper functions for creating GiglResourceConfig configurations - - -def _create_resource_config_with_trainer_graph_store() -> GiglResourceConfigWrapper: - """Create a GiglResourceConfig with VertexAiGraphStoreConfig for trainer.""" +def _create_resource_config_with_both_graph_stores() -> GiglResourceConfigWrapper: + """Create a GiglResourceConfig with VertexAiGraphStoreConfig for both trainer and inferencer.""" config = gigl_resource_config_pb2.GiglResourceConfig() _create_shared_resource_config(config) @@ -137,15 +102,15 @@ def _create_resource_config_with_trainer_graph_store() -> GiglResourceConfigWrap config.trainer_resource_config.vertex_ai_graph_store_trainer_config.CopyFrom( _create_vertex_ai_graph_store_config() ) - # Inferencer with standard config - config.inferencer_resource_config.vertex_ai_inferencer_config.CopyFrom( - _create_vertex_ai_resource_config() + # Inferencer with VertexAiGraphStoreConfig + config.inferencer_resource_config.vertex_ai_graph_store_inferencer_config.CopyFrom( + _create_vertex_ai_graph_store_config() ) return GiglResourceConfigWrapper(resource_config=config) -def _create_resource_config_without_trainer_graph_store() -> GiglResourceConfigWrapper: - """Create a GiglResourceConfig without VertexAiGraphStoreConfig for trainer.""" +def _create_resource_config_without_graph_stores() -> GiglResourceConfigWrapper: + """Create a GiglResourceConfig without VertexAiGraphStoreConfig for trainer or inferencer.""" config = gigl_resource_config_pb2.GiglResourceConfig() _create_shared_resource_config(config) @@ -160,60 +125,13 @@ def _create_resource_config_without_trainer_graph_store() -> GiglResourceConfigW return GiglResourceConfigWrapper(resource_config=config) -def _create_resource_config_with_inferencer_graph_store() -> GiglResourceConfigWrapper: - """Create a GiglResourceConfig with VertexAiGraphStoreConfig for inferencer.""" - config = gigl_resource_config_pb2.GiglResourceConfig() - _create_shared_resource_config(config) - - # Trainer with standard config - config.trainer_resource_config.vertex_ai_trainer_config.CopyFrom( - _create_vertex_ai_resource_config() - ) - # Inferencer with VertexAiGraphStoreConfig - config.inferencer_resource_config.vertex_ai_graph_store_inferencer_config.CopyFrom( - _create_vertex_ai_graph_store_config() - ) - return GiglResourceConfigWrapper(resource_config=config) - - -def _create_resource_config_without_inferencer_graph_store() -> ( - GiglResourceConfigWrapper -): - """Create a GiglResourceConfig without VertexAiGraphStoreConfig for inferencer.""" - return _create_resource_config_without_trainer_graph_store() - - -def _create_resource_config_with_both_graph_stores() -> GiglResourceConfigWrapper: - """Create a GiglResourceConfig with VertexAiGraphStoreConfig for both trainer and inferencer.""" - config = gigl_resource_config_pb2.GiglResourceConfig() - _create_shared_resource_config(config) - - # Trainer with VertexAiGraphStoreConfig - config.trainer_resource_config.vertex_ai_graph_store_trainer_config.CopyFrom( - _create_vertex_ai_graph_store_config() - ) - # Inferencer with VertexAiGraphStoreConfig - config.inferencer_resource_config.vertex_ai_graph_store_inferencer_config.CopyFrom( - _create_vertex_ai_graph_store_config() - ) - return GiglResourceConfigWrapper(resource_config=config) - - -def _create_resource_config_without_graph_stores() -> GiglResourceConfigWrapper: - """Create a GiglResourceConfig without VertexAiGraphStoreConfig for both trainer and inferencer.""" - return _create_resource_config_without_trainer_graph_store() - - -# Test Classes - - class TestTrainerGraphStoreCompatibility(unittest.TestCase): """Test suite for trainer graph store compatibility checks.""" def test_both_have_trainer_graph_store(self): """Test that both configs having trainer graph store passes validation.""" - gbml_config = _create_gbml_config_with_trainer_graph_store() - resource_config = _create_resource_config_with_trainer_graph_store() + gbml_config = _create_gbml_config_with_both_graph_stores() + resource_config = _create_resource_config_with_both_graph_stores() # Should not raise any exception check_trainer_graph_store_compatibility( gbml_config_pb_wrapper=gbml_config, @@ -222,8 +140,8 @@ def test_both_have_trainer_graph_store(self): def test_neither_has_trainer_graph_store(self): """Test that neither config having trainer graph store passes validation.""" - gbml_config = _create_gbml_config_without_trainer_graph_store() - resource_config = _create_resource_config_without_trainer_graph_store() + gbml_config = _create_gbml_config_without_graph_stores() + resource_config = _create_resource_config_without_graph_stores() # Should not raise any exception check_trainer_graph_store_compatibility( gbml_config_pb_wrapper=gbml_config, @@ -232,8 +150,8 @@ def test_neither_has_trainer_graph_store(self): def test_template_has_trainer_graph_store_resource_does_not(self): """Test that template having graph store but resource not raises an assertion error.""" - gbml_config = _create_gbml_config_with_trainer_graph_store() - resource_config = _create_resource_config_without_trainer_graph_store() + gbml_config = _create_gbml_config_with_both_graph_stores() + resource_config = _create_resource_config_without_graph_stores() with self.assertRaises(AssertionError): check_trainer_graph_store_compatibility( gbml_config_pb_wrapper=gbml_config, @@ -242,8 +160,8 @@ def test_template_has_trainer_graph_store_resource_does_not(self): def test_resource_has_trainer_graph_store_template_does_not(self): """Test that resource having graph store but template not raises an assertion error.""" - gbml_config = _create_gbml_config_without_trainer_graph_store() - resource_config = _create_resource_config_with_trainer_graph_store() + gbml_config = _create_gbml_config_without_graph_stores() + resource_config = _create_resource_config_with_both_graph_stores() with self.assertRaises(AssertionError): check_trainer_graph_store_compatibility( gbml_config_pb_wrapper=gbml_config, @@ -256,8 +174,8 @@ class TestInferencerGraphStoreCompatibility(unittest.TestCase): def test_both_have_inferencer_graph_store(self): """Test that both configs having inferencer graph store passes validation.""" - gbml_config = _create_gbml_config_with_inferencer_graph_store() - resource_config = _create_resource_config_with_inferencer_graph_store() + gbml_config = _create_gbml_config_with_both_graph_stores() + resource_config = _create_resource_config_with_both_graph_stores() # Should not raise any exception check_inferencer_graph_store_compatibility( gbml_config_pb_wrapper=gbml_config, @@ -266,8 +184,8 @@ def test_both_have_inferencer_graph_store(self): def test_neither_has_inferencer_graph_store(self): """Test that neither config having inferencer graph store passes validation.""" - gbml_config = _create_gbml_config_without_inferencer_graph_store() - resource_config = _create_resource_config_without_inferencer_graph_store() + gbml_config = _create_gbml_config_without_graph_stores() + resource_config = _create_resource_config_without_graph_stores() # Should not raise any exception check_inferencer_graph_store_compatibility( gbml_config_pb_wrapper=gbml_config, @@ -276,8 +194,8 @@ def test_neither_has_inferencer_graph_store(self): def test_template_has_inferencer_graph_store_resource_does_not(self): """Test that template having graph store but resource not raises an assertion error.""" - gbml_config = _create_gbml_config_with_inferencer_graph_store() - resource_config = _create_resource_config_without_inferencer_graph_store() + gbml_config = _create_gbml_config_with_both_graph_stores() + resource_config = _create_resource_config_without_graph_stores() with self.assertRaises(AssertionError): check_inferencer_graph_store_compatibility( gbml_config_pb_wrapper=gbml_config, @@ -286,110 +204,14 @@ def test_template_has_inferencer_graph_store_resource_does_not(self): def test_resource_has_inferencer_graph_store_template_does_not(self): """Test that resource having graph store but template not raises an assertion error.""" - gbml_config = _create_gbml_config_without_inferencer_graph_store() - resource_config = _create_resource_config_with_inferencer_graph_store() - with self.assertRaises(AssertionError): - check_inferencer_graph_store_compatibility( - gbml_config_pb_wrapper=gbml_config, - resource_config_wrapper=resource_config, - ) - - -class TestMixedGraphStoreConfigurations(unittest.TestCase): - """Test suite for mixed graph store configuration scenarios.""" - - def test_both_have_all_graph_stores(self): - """Test that both configs having all graph stores passes validation.""" - gbml_config = _create_gbml_config_with_both_graph_stores() - resource_config = _create_resource_config_with_both_graph_stores() - # Should not raise any exception for trainer - check_trainer_graph_store_compatibility( - gbml_config_pb_wrapper=gbml_config, - resource_config_wrapper=resource_config, - ) - # Should not raise any exception for inferencer - check_inferencer_graph_store_compatibility( - gbml_config_pb_wrapper=gbml_config, - resource_config_wrapper=resource_config, - ) - - def test_neither_has_any_graph_stores(self): - """Test that neither config having any graph stores passes validation.""" gbml_config = _create_gbml_config_without_graph_stores() - resource_config = _create_resource_config_without_graph_stores() - # Should not raise any exception for trainer - check_trainer_graph_store_compatibility( - gbml_config_pb_wrapper=gbml_config, - resource_config_wrapper=resource_config, - ) - # Should not raise any exception for inferencer - check_inferencer_graph_store_compatibility( - gbml_config_pb_wrapper=gbml_config, - resource_config_wrapper=resource_config, - ) - - def test_trainer_graph_store_only_compatible(self): - """Test trainer graph store only configuration is compatible.""" - gbml_config = _create_gbml_config_with_trainer_graph_store() - resource_config = _create_resource_config_with_trainer_graph_store() - # Should not raise any exception for trainer - check_trainer_graph_store_compatibility( - gbml_config_pb_wrapper=gbml_config, - resource_config_wrapper=resource_config, - ) - # Should not raise any exception for inferencer (neither has it) - check_inferencer_graph_store_compatibility( - gbml_config_pb_wrapper=gbml_config, - resource_config_wrapper=resource_config, - ) - - def test_inferencer_graph_store_only_compatible(self): - """Test inferencer graph store only configuration is compatible.""" - gbml_config = _create_gbml_config_with_inferencer_graph_store() - resource_config = _create_resource_config_with_inferencer_graph_store() - # Should not raise any exception for trainer (neither has it) - check_trainer_graph_store_compatibility( - gbml_config_pb_wrapper=gbml_config, - resource_config_wrapper=resource_config, - ) - # Should not raise any exception for inferencer - check_inferencer_graph_store_compatibility( - gbml_config_pb_wrapper=gbml_config, - resource_config_wrapper=resource_config, - ) - - def test_template_has_both_resource_has_trainer_only(self): - """Test that template having both but resource having only trainer raises an error for inferencer.""" - gbml_config = _create_gbml_config_with_both_graph_stores() - resource_config = _create_resource_config_with_trainer_graph_store() - # Should not raise any exception for trainer - check_trainer_graph_store_compatibility( - gbml_config_pb_wrapper=gbml_config, - resource_config_wrapper=resource_config, - ) - # Should raise an assertion error for inferencer + resource_config = _create_resource_config_with_both_graph_stores() with self.assertRaises(AssertionError): check_inferencer_graph_store_compatibility( gbml_config_pb_wrapper=gbml_config, resource_config_wrapper=resource_config, ) - def test_template_has_both_resource_has_inferencer_only(self): - """Test that template having both but resource having only inferencer raises an error for trainer.""" - gbml_config = _create_gbml_config_with_both_graph_stores() - resource_config = _create_resource_config_with_inferencer_graph_store() - # Should raise an assertion error for trainer - with self.assertRaises(AssertionError): - check_trainer_graph_store_compatibility( - gbml_config_pb_wrapper=gbml_config, - resource_config_wrapper=resource_config, - ) - # Should not raise any exception for inferencer - check_inferencer_graph_store_compatibility( - gbml_config_pb_wrapper=gbml_config, - resource_config_wrapper=resource_config, - ) - if __name__ == "__main__": unittest.main() From b3019b41c9c3d28f5d929bafb48ffb70e234876c Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Fri, 30 Jan 2026 18:12:34 +0000 Subject: [PATCH 5/6] cleanup --- .../lib/resource_config_checks_test.py | 38 ++++++------------- 1 file changed, 11 insertions(+), 27 deletions(-) diff --git a/tests/unit/src/validation/lib/resource_config_checks_test.py b/tests/unit/src/validation/lib/resource_config_checks_test.py index eff49aaab..d136b85a8 100644 --- a/tests/unit/src/validation/lib/resource_config_checks_test.py +++ b/tests/unit/src/validation/lib/resource_config_checks_test.py @@ -756,34 +756,20 @@ def test_valid_vertex_ai_graph_store_config(self): # Helper functions for creating GbmlConfig configurations -def _create_gbml_config_with_trainer_graph_store( +def _create_gbml_config_with_both_graph_stores( storage_command: str = "python -m gigl.distributed.graph_store.storage_main", ) -> GbmlConfigPbWrapper: - """Create a GbmlConfig with graph_store_storage_config set for trainer.""" + """Create a GbmlConfig with graph_store_storage_config set for both trainer and inferencer.""" gbml_config = gbml_config_pb2.GbmlConfig() gbml_config.trainer_config.graph_store_storage_config.command = storage_command - return GbmlConfigPbWrapper(gbml_config_pb=gbml_config) - - -def _create_gbml_config_without_trainer_graph_store() -> GbmlConfigPbWrapper: - """Create a GbmlConfig without graph_store_storage_config for trainer.""" - gbml_config = gbml_config_pb2.GbmlConfig() - gbml_config.trainer_config.trainer_args["some_arg"] = "some_value" - return GbmlConfigPbWrapper(gbml_config_pb=gbml_config) - - -def _create_gbml_config_with_inferencer_graph_store( - storage_command: str = "python -m gigl.distributed.graph_store.storage_main", -) -> GbmlConfigPbWrapper: - """Create a GbmlConfig with graph_store_storage_config set for inferencer.""" - gbml_config = gbml_config_pb2.GbmlConfig() gbml_config.inferencer_config.graph_store_storage_config.command = storage_command return GbmlConfigPbWrapper(gbml_config_pb=gbml_config) -def _create_gbml_config_without_inferencer_graph_store() -> GbmlConfigPbWrapper: - """Create a GbmlConfig without graph_store_storage_config for inferencer.""" +def _create_gbml_config_without_graph_stores() -> GbmlConfigPbWrapper: + """Create a GbmlConfig without graph_store_storage_config for trainer or inferencer.""" gbml_config = gbml_config_pb2.GbmlConfig() + gbml_config.trainer_config.trainer_args["some_arg"] = "some_value" gbml_config.inferencer_config.inferencer_args["some_arg"] = "some_value" return GbmlConfigPbWrapper(gbml_config_pb=gbml_config) @@ -793,19 +779,19 @@ class TestTrainerGraphStoreStorageCommand(unittest.TestCase): def test_valid_storage_command(self): """Test that a valid storage_command passes validation.""" - gbml_config = _create_gbml_config_with_trainer_graph_store() + gbml_config = _create_gbml_config_with_both_graph_stores() # Should not raise any exception check_if_trainer_graph_store_storage_command_valid(gbml_config) def test_missing_storage_command(self): """Test that missing storage_command raises an assertion error.""" - gbml_config = _create_gbml_config_with_trainer_graph_store(storage_command="") + gbml_config = _create_gbml_config_with_both_graph_stores(storage_command="") with self.assertRaises(AssertionError): check_if_trainer_graph_store_storage_command_valid(gbml_config) def test_no_graph_store_config(self): """Test that no graph store config passes validation (nothing to check).""" - gbml_config = _create_gbml_config_without_trainer_graph_store() + gbml_config = _create_gbml_config_without_graph_stores() # Should not raise any exception - no graph store means nothing to validate check_if_trainer_graph_store_storage_command_valid(gbml_config) @@ -815,21 +801,19 @@ class TestInferencerGraphStoreStorageCommand(unittest.TestCase): def test_valid_storage_command(self): """Test that a valid storage_command passes validation.""" - gbml_config = _create_gbml_config_with_inferencer_graph_store() + gbml_config = _create_gbml_config_with_both_graph_stores() # Should not raise any exception check_if_inferencer_graph_store_storage_command_valid(gbml_config) def test_missing_storage_command(self): """Test that missing storage_command raises an assertion error.""" - gbml_config = _create_gbml_config_with_inferencer_graph_store( - storage_command="" - ) + gbml_config = _create_gbml_config_with_both_graph_stores(storage_command="") with self.assertRaises(AssertionError): check_if_inferencer_graph_store_storage_command_valid(gbml_config) def test_no_graph_store_config(self): """Test that no graph store config passes validation (nothing to check).""" - gbml_config = _create_gbml_config_without_inferencer_graph_store() + gbml_config = _create_gbml_config_without_graph_stores() # Should not raise any exception - no graph store means nothing to validate check_if_inferencer_graph_store_storage_command_valid(gbml_config) From a11b6be798e2bc14a3c36e108ef269800163886b Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Tue, 3 Feb 2026 05:03:47 +0000 Subject: [PATCH 6/6] away from xor --- .../libs/gbml_and_resource_config_compatibility_checks.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/gigl/src/validation_check/libs/gbml_and_resource_config_compatibility_checks.py b/gigl/src/validation_check/libs/gbml_and_resource_config_compatibility_checks.py index 84dd75831..fc12d1939 100644 --- a/gigl/src/validation_check/libs/gbml_and_resource_config_compatibility_checks.py +++ b/gigl/src/validation_check/libs/gbml_and_resource_config_compatibility_checks.py @@ -96,7 +96,7 @@ def check_trainer_graph_store_compatibility( resource_config_wrapper, "trainer" ) - if gbml_has_graph_store ^ resource_has_graph_store: + if gbml_has_graph_store != resource_has_graph_store: raise AssertionError( f"If one of GbmlConfig.trainer_config.graph_store_storage_config or GiglResourceConfig.trainer_resource_config is set, the other must also be set. GbmlConfig.trainer_config.graph_store_storage_config is set: {gbml_has_graph_store}, GiglResourceConfig.trainer_resource_config is set: {resource_has_graph_store}." ) @@ -131,7 +131,7 @@ def check_inferencer_graph_store_compatibility( resource_config_wrapper, "inferencer" ) - if gbml_has_graph_store ^ resource_has_graph_store: + if gbml_has_graph_store != resource_has_graph_store: raise AssertionError( f"If one of GbmlConfig.inferencer_config.graph_store_storage_config or GiglResourceConfig.inferencer_resource_config is set, the other must also be set. GbmlConfig.inferencer_config.graph_store_storage_config is set: {gbml_has_graph_store}, GiglResourceConfig.inferencer_resource_config is set: {resource_has_graph_store}." )