diff --git a/deployment/configs/e2e/preprocessor.yaml b/deployment/configs/e2e/preprocessor.yaml new file mode 100644 index 000000000..4f5c20cb4 --- /dev/null +++ b/deployment/configs/e2e/preprocessor.yaml @@ -0,0 +1,11 @@ +# Shared preprocessing capacity for GiGL's end-to-end tests. +edge_preprocessor_config: + num_workers: 1 + max_num_workers: 128 + machine_type: "n2d-highmem-32" + disk_size_gb: 300 +node_preprocessor_config: + num_workers: 1 + max_num_workers: 128 + machine_type: "n2d-highmem-64" + disk_size_gb: 300 diff --git a/deployment/configs/e2e/shared.yaml b/deployment/configs/e2e/shared.yaml new file mode 100644 index 000000000..72d7c61c2 --- /dev/null +++ b/deployment/configs/e2e/shared.yaml @@ -0,0 +1,15 @@ +# Shared infrastructure for GiGL's end-to-end test resource configs. +resource_labels: + # COMPONENT is filled by GiGL with the active pipeline component. + cost_resource_group_tag: dev_experiments_COMPONENT + cost_resource_group: gigl_platform +common_compute_config: + project: "external-snap-ci-github-gigl" + region: "us-central1" + temp_assets_bucket: "gs://gigl-cicd-temp" + temp_regional_assets_bucket: "gs://gigl-cicd-temp" + perm_assets_bucket: "gs://gigl-cicd-perm" + temp_assets_bq_dataset_name: "gigl_temp_assets" + embedding_bq_dataset_name: "gigl_embeddings" + gcp_service_account_email: "untrusted-external-github-gigl@external-snap-ci-github-gigl.iam.gserviceaccount.com" + dataflow_runner: "DataflowRunner" diff --git a/deployment/configs/e2e_cicd_resource_config.yaml b/deployment/configs/e2e_cicd_resource_config.yaml index 21fe6ab0a..d689fee9f 100644 --- a/deployment/configs/e2e_cicd_resource_config.yaml +++ b/deployment/configs/e2e_cicd_resource_config.yaml @@ -1,38 +1,9 @@ # GiglResourceConfig for our e2e intergration tests. -shared_resource_config: - # Resource labels are just compute labels that should be attached to all compute resources spun up by GiGL. - # So a practitioner can have a more fine grained understanding of resource utilization and cost of the resources in their GCP billing. - # Read more here: https://cloud.google.com/compute/docs/labeling-resources#what-are-labels - resource_labels: - # We have a 63 character limit for cost_resource_group_tag. +defaults: + - e2e/shared@shared_resource_config + - e2e/preprocessor@preprocessor_config + - _self_ - # COMPONENT is one of {pre|sgs|spl|tra|inf|pos} standing for: - # {Preprocessor | Subgraph Sampler | Split Generator | Trainer | Inference - # | Post Processor} so we can get more accurate cost measurements - # of each component. This will be automatically filled in code. - cost_resource_group_tag: dev_experiments_COMPONENT - cost_resource_group: gigl_platform - common_compute_config: - project: "external-snap-ci-github-gigl" - region: "us-central1" - temp_assets_bucket: "gs://gigl-cicd-temp" - temp_regional_assets_bucket: "gs://gigl-cicd-temp" - perm_assets_bucket: "gs://gigl-cicd-perm" - temp_assets_bq_dataset_name: "gigl_temp_assets" - embedding_bq_dataset_name: "gigl_embeddings" - gcp_service_account_email: "untrusted-external-github-gigl@external-snap-ci-github-gigl.iam.gserviceaccount.com" - dataflow_runner: "DataflowRunner" -preprocessor_config: - edge_preprocessor_config: - num_workers: 1 - max_num_workers: 128 - machine_type: "n2d-highmem-32" - disk_size_gb: 300 - node_preprocessor_config: - num_workers: 1 - max_num_workers: 128 - machine_type: "n2d-highmem-64" - disk_size_gb: 300 subgraph_sampler_config: machine_type: "n2d-highmem-16" num_local_ssds: 2 diff --git a/deployment/configs/e2e_glt_gs_resource_config.yaml b/deployment/configs/e2e_glt_gs_resource_config.yaml index 4839c8f48..7abfb83ff 100644 --- a/deployment/configs/e2e_glt_gs_resource_config.yaml +++ b/deployment/configs/e2e_glt_gs_resource_config.yaml @@ -1,31 +1,11 @@ # Diffs from e2e_glt_resource_config.yaml # - Swap vertex_ai_inferencer_config for vertex_ai_graph_store_inferencer_config # - Swap vertex_ai_trainer_config for vertex_ai_graph_store_trainer_config -shared_resource_config: - resource_labels: - cost_resource_group_tag: dev_experiments_COMPONENT - cost_resource_group: gigl_platform - common_compute_config: - project: "external-snap-ci-github-gigl" - region: "us-central1" - temp_assets_bucket: "gs://gigl-cicd-temp" - temp_regional_assets_bucket: "gs://gigl-cicd-temp" - perm_assets_bucket: "gs://gigl-cicd-perm" - temp_assets_bq_dataset_name: "gigl_temp_assets" - embedding_bq_dataset_name: "gigl_embeddings" - gcp_service_account_email: "untrusted-external-github-gigl@external-snap-ci-github-gigl.iam.gserviceaccount.com" - dataflow_runner: "DataflowRunner" -preprocessor_config: - edge_preprocessor_config: - num_workers: 1 - max_num_workers: 128 - machine_type: "n2d-highmem-32" - disk_size_gb: 300 - node_preprocessor_config: - num_workers: 1 - max_num_workers: 128 - machine_type: "n2d-highmem-64" - disk_size_gb: 300 +defaults: + - e2e/shared@shared_resource_config + - e2e/preprocessor@preprocessor_config + - _self_ + trainer_resource_config: vertex_ai_graph_store_trainer_config: graph_store_pool: diff --git a/deployment/configs/e2e_glt_resource_config.yaml b/deployment/configs/e2e_glt_resource_config.yaml index 540c4ace1..7b08a37c4 100644 --- a/deployment/configs/e2e_glt_resource_config.yaml +++ b/deployment/configs/e2e_glt_resource_config.yaml @@ -1,28 +1,8 @@ -shared_resource_config: - resource_labels: - cost_resource_group_tag: dev_experiments_COMPONENT - cost_resource_group: gigl_platform - common_compute_config: - project: "external-snap-ci-github-gigl" - region: "us-central1" - temp_assets_bucket: "gs://gigl-cicd-temp" - temp_regional_assets_bucket: "gs://gigl-cicd-temp" - perm_assets_bucket: "gs://gigl-cicd-perm" - temp_assets_bq_dataset_name: "gigl_temp_assets" - embedding_bq_dataset_name: "gigl_embeddings" - gcp_service_account_email: "untrusted-external-github-gigl@external-snap-ci-github-gigl.iam.gserviceaccount.com" - dataflow_runner: "DataflowRunner" -preprocessor_config: - edge_preprocessor_config: - num_workers: 1 - max_num_workers: 128 - machine_type: "n2d-highmem-32" - disk_size_gb: 300 - node_preprocessor_config: - num_workers: 1 - max_num_workers: 128 - machine_type: "n2d-highmem-64" - disk_size_gb: 300 +defaults: + - e2e/shared@shared_resource_config + - e2e/preprocessor@preprocessor_config + - _self_ + trainer_resource_config: vertex_ai_trainer_config: machine_type: n1-highmem-32 diff --git a/docs/user_guide/config_guides/hydra_composition.md b/docs/user_guide/config_guides/hydra_composition.md new file mode 100644 index 000000000..9ef4002d4 --- /dev/null +++ b/docs/user_guide/config_guides/hydra_composition.md @@ -0,0 +1,120 @@ +# Composing task and resource configs with Hydra + +Task and resource configs can use [Hydra Defaults Lists](https://hydra.cc/docs/advanced/defaults_list/) by adding a +top-level `defaults` list. ConfigValidator composes every source config before publishing plain YAML snapshots. +`ProtoUtils.read_proto_from_yaml` remains the URI-agnostic, non-composing reader; `ProtoUtils.compose_proto_from_yaml` +is the URI-agnostic composition API. + +## Config root + +The primary file's parent directory is its Hydra config root. GiGL does not prevalidate filename extensions; Hydra 1.3 +resolves local config names with its `.yaml` convention. Remote primaries are staged with a `.yaml` extension before +composition. + +```text +configs/ +├── task.yaml +├── resource.yaml +├── shared/ +│ └── directed.yaml +└── compute/ + └── local.yaml +``` + +Keep primary files directly in the bundle root. GiGL does not search for a repository root or a directory named +`configs`. Local configs retain their parent as the config root so they can select sibling fragments. GCS and HTTP +primaries are downloaded to a temporary local file before composition; relative Defaults List entries are not downloaded +with them. Remote primaries can opt into an existing local config root as described below. + +## Remote primary with local shared configs + +GCS and HTTPS primaries can select shared configs installed in the runtime image by adding a Hydra search path: + +```yaml +hydra: + searchpath: + - pkg://my_project.configs + +defaults: + - shared@_global_: common + - _self_ +``` + +Prefer `pkg://` roots so the same config works across machines and containers. The referenced package and its YAML files +must be installed in every environment that composes the primary. Hydra removes its own `hydra` metadata from the +composed mapping before GiGL parses the protobuf. + +GiGL does not download Defaults List entries relative to a GCS prefix or HTTPS location. Without an explicit search +path, a remote primary is composed as a standalone file. + +## Task config example + +```yaml +# task.yaml +defaults: + - shared@sharedConfig: directed + - _self_ +``` + +```yaml +# shared/directed.yaml +isGraphDirected: true +``` + +The package target after `@` places the selected group at the corresponding protobuf field. A fragment containing fields +for the whole protobuf can use `@_global_`. + +Include `_self_` explicitly so it is clear whether values in the primary override group values or are overridden by +them. + +## Resource config example + +```yaml +# resource.yaml +defaults: + - compute@shared_resource_config: local + - _self_ +``` + +```yaml +# compute/local.yaml +common_compute_config: + project: example-project + region: us-central1 + temp_regional_assets_bucket: gs://example-bucket +``` + +KFP submission composes a local resource source once to select the project, region, service account, and staging bucket. +ConfigValidator composes it again inside the pipeline, and that validation-time result becomes the resource config used +by downstream components. Resolvers such as `now`, `git_hash`, and `oc.env` are supported; environment-based resolvers +must be available wherever the source config is composed. + +## Pipeline behavior + +ConfigValidator composes both primary configs, writes fully resolved plain-protobuf YAML snapshots, initializes its +runtime from those exact snapshots, and validates them before publishing its outputs. Every downstream pipeline +component receives those snapshot URIs. ConfigValidator caching is disabled so changing a selected fragment cannot reuse +stale resolved outputs. + +The final composed mapping must still be a valid `GbmlConfig` or `GiglResourceConfig`. Protobuf parsing remains the +schema and type validation boundary. + +## Repository example + +GiGL's three E2E resource configs share their infrastructure and preprocessing sections: + +- [`e2e_cicd_resource_config.yaml`](../../../deployment/configs/e2e_cicd_resource_config.yaml) +- [`e2e_glt_resource_config.yaml`](../../../deployment/configs/e2e_glt_resource_config.yaml) +- [`e2e_glt_gs_resource_config.yaml`](../../../deployment/configs/e2e_glt_gs_resource_config.yaml) + +Each primary selects [`e2e/shared.yaml`](../../../deployment/configs/e2e/shared.yaml) and +[`e2e/preprocessor.yaml`](../../../deployment/configs/e2e/preprocessor.yaml), then defines only its pipeline-specific +resources. Their names intentionally omit `resource_config` so repository validation does not mistake these partial +fragments for complete resource configs. The unit-test resource config remains standalone because its buckets, datasets, +and runner intentionally differ. + +## Boundaries + +- GiGL does not consume Hydra command-line overrides, multirun, launchers, or output-directory behavior. +- GiGL does not automatically download config fragments next to a GCS or HTTPS primary. +- Treat config bundle write access as trusted access. GiGL configs can reference importable classes and commands. diff --git a/docs/user_guide/config_guides/resource_config_guide.md b/docs/user_guide/config_guides/resource_config_guide.md index 900b57af3..76e6b3c35 100644 --- a/docs/user_guide/config_guides/resource_config_guide.md +++ b/docs/user_guide/config_guides/resource_config_guide.md @@ -4,6 +4,9 @@ GiGL Resource Config is a yaml file that is passed in at runtime and specifies t configurations for each component in the GiGL. The proto definition for GiGL Resource Config can be seen [here](https://github.com/Snapchat/GiGL/blob/main/proto/snapchat/research/gbml/gigl_resource_config.proto) +Resource configs can be split into Hydra config groups. See +[composing task and resource configs with Hydra](./hydra_composition.md). + ## Prerequisites If you don't have cloud assets already setup i.e. a GCP project. See [guide](../getting_started/cloud_setup_guide.md) diff --git a/docs/user_guide/config_guides/task_config_guide.md b/docs/user_guide/config_guides/task_config_guide.md index 0fd701282..93fe1ed3d 100644 --- a/docs/user_guide/config_guides/task_config_guide.md +++ b/docs/user_guide/config_guides/task_config_guide.md @@ -3,6 +3,9 @@ The task config specifies task-related configurations - guiding the behavior of components according to the needs of your machine learning task. +Task configs can be split into Hydra config groups. See +[composing task and resource configs with Hydra](./hydra_composition.md). + Whenever we say "task config" we are talking about an instance off {py:class}`snapchat.research.gbml.gbml_config_pb2.GbmlConfig` diff --git a/gigl/common/utils/hydra_config.py b/gigl/common/utils/hydra_config.py new file mode 100644 index 000000000..d9ef817f7 --- /dev/null +++ b/gigl/common/utils/hydra_config.py @@ -0,0 +1,72 @@ +"""Hydra composition support for GiGL protobuf YAML configs.""" + +from __future__ import annotations + +import os +import threading +from pathlib import Path +from typing import Any, cast + +from hydra import compose, initialize_config_dir +from hydra.core.global_hydra import GlobalHydra +from omegaconf import OmegaConf + +from gigl.common import LocalUri +from gigl.common.omegaconf_resolvers import now_resolver, register_resolvers + +_COMPOSE_LOCK = threading.RLock() + + +def compose_yaml_config(uri: LocalUri) -> dict[str, Any]: + """Compose a YAML config with Hydra using its parent as the config root. + + Args: + uri: Primary YAML config URI. + + Returns: + A fully composed and resolved mapping. + + Raises: + ValueError: If the result is not a mapping. + RuntimeError: If another Hydra application owns the global context. + """ + primary_name = uri.get_basename() + config_name = primary_name.rsplit(".", 1)[0] + + config_root = Path(uri.uri).absolute().parent + + with _COMPOSE_LOCK: + if GlobalHydra.instance().is_initialized(): + raise RuntimeError( + "GiGL cannot compose a config while another Hydra context is active." + ) + try: + with initialize_config_dir( + config_dir=os.fspath(config_root), + job_name="gigl_config", + version_base="1.3", + ): + # Hydra installs a one-argument ``now`` resolver during + # initialization. GiGL's resolver is a backward-compatible + # superset that also supports offsets. + register_resolvers() + OmegaConf.register_new_resolver( + "now", + now_resolver, + replace=True, + ) + composed = compose(config_name=config_name, overrides=[]) + resolved = OmegaConf.to_container(composed, resolve=True) + finally: + OmegaConf.register_new_resolver( + "now", + now_resolver, + replace=True, + ) + + if not isinstance(resolved, dict): + raise ValueError( + f"Hydra config {config_root / config_name} resolved to " + f"{type(resolved).__name__}, expected a mapping." + ) + return cast(dict[str, Any], resolved) diff --git a/gigl/common/utils/proto_utils.py b/gigl/common/utils/proto_utils.py index e61258756..4f3b2fd33 100644 --- a/gigl/common/utils/proto_utils.py +++ b/gigl/common/utils/proto_utils.py @@ -9,6 +9,7 @@ from gigl.common import LocalUri, Uri from gigl.common.logger import Logger from gigl.common.omegaconf_resolvers import register_resolvers +from gigl.common.utils.hydra_config import compose_yaml_config from gigl.src.common.utils.file_loader import FileLoader logger = Logger() @@ -16,6 +17,19 @@ T = TypeVar("T", bound=message.Message) +def proto_to_yaml(proto: message.Message) -> str: + """Serialize a protobuf message to canonical YAML. + + Args: + proto: Protobuf message to serialize. + + Returns: + YAML containing the protobuf JSON representation. + """ + proto_dict = MessageToDict(message=proto) + return yaml.safe_dump(proto_dict, default_flow_style=False, sort_keys=True) + + class ProtoUtils: def __init__(self, project: Optional[str] = None) -> None: self.__file_loader = FileLoader(project=project) @@ -36,6 +50,27 @@ def read_proto_from_yaml(self, uri: Uri, proto_cls: Type[T]) -> T: proto = ParseDict(js_dict=cast(dict, obj_dict), message=proto_cls()) return proto + def compose_proto_from_yaml(self, uri: Uri, proto_cls: Type[T]) -> T: + """Compose a YAML config with Hydra and parse it as a protobuf. + + Remote URIs are downloaded as a single primary config; relative Defaults + List entries are not fetched. + """ + if isinstance(uri, LocalUri): + # Compose from the source path so its parent remains Hydra's config + # root and relative Defaults List entries resolve as expected. + obj_dict = compose_yaml_config(uri=uri) + else: + with NamedTemporaryFile(suffix=".yaml") as temp_file: + local_uri = LocalUri(temp_file.name) + self.__file_loader.load_file( + file_uri_src=uri, + file_uri_dst=local_uri, + ) + obj_dict = compose_yaml_config(uri=local_uri) + proto = ParseDict(js_dict=cast(dict, obj_dict), message=proto_cls()) + return proto + def read_proto_from_binary(self, uri: Uri, proto_cls: Type[T]) -> T: tfh = self.__file_loader.load_to_temp_file(file_uri_src=uri, delete=False) with open(tfh.name, "rb") as file: diff --git a/gigl/orchestration/kubeflow/components/config_validator/component.yaml b/gigl/orchestration/kubeflow/components/config_validator/component.yaml index 7c4ec057e..fc1963c4a 100644 --- a/gigl/orchestration/kubeflow/components/config_validator/component.yaml +++ b/gigl/orchestration/kubeflow/components/config_validator/component.yaml @@ -9,6 +9,9 @@ inputs: - {name: cpu_docker_uri, type: String, description: "Uri to dockerized source code compiled for cpu at runtime"} - {name: cuda_docker_uri, type: String, description: "Uri to dockerized source code compiled for gpu at runtime"} outputs: +- {name: resolved_task_config_uri, type: String, description: 'Resolved task config URI'} +- {name: resolved_resource_config_uri, type: String, description: 'Resolved resource config URI'} +- {name: should_use_glt_backend, type: Boolean, description: 'Whether to use the GLT backend'} implementation: container: @@ -21,5 +24,8 @@ implementation: --resource_config_uri, {inputValue: resource_config_uri}, --stop_after, {inputValue: stop_after}, --cpu_docker_uri, {inputValue: cpu_docker_uri}, - --cuda_docker_uri, {inputValue: cuda_docker_uri} + --cuda_docker_uri, {inputValue: cuda_docker_uri}, + --output_file_path_resolved_task_config_uri, {outputPath: resolved_task_config_uri}, + --output_file_path_resolved_resource_config_uri, {outputPath: resolved_resource_config_uri}, + --output_file_path_should_use_glt_backend, {outputPath: should_use_glt_backend} ] diff --git a/gigl/orchestration/kubeflow/kfp_orchestrator.py b/gigl/orchestration/kubeflow/kfp_orchestrator.py index 374347ac1..56638d2bd 100644 --- a/gigl/orchestration/kubeflow/kfp_orchestrator.py +++ b/gigl/orchestration/kubeflow/kfp_orchestrator.py @@ -12,15 +12,19 @@ from gigl.common.logger import Logger from gigl.common.services.vertex_ai import VertexAIService from gigl.common.types.resource_config import CommonPipelineComponentConfigs -from gigl.env.pipelines_config import get_resource_config +from gigl.common.utils.proto_utils import ProtoUtils from gigl.orchestration.kubeflow.kfp_pipeline import generate_pipeline from gigl.src.common.constants.components import GiGLComponents from gigl.src.common.types import AppliedTaskIdentifier +from gigl.src.common.types.pb_wrappers.gigl_resource_config import ( + GiglResourceConfigWrapper, +) from gigl.src.common.utils.file_loader import FileLoader from gigl.src.common.utils.time import current_formatted_datetime from gigl.src.validation_check.libs.name_checks import ( check_if_kfp_pipeline_job_name_valid, ) +from snapchat.research.gbml.gigl_resource_config_pb2 import GiglResourceConfig logger = Logger() @@ -143,6 +147,12 @@ def run( ) logger.info(f"Skipping pipeline compilation; will use {compiled_pipeline_path}") + resource_config = GiglResourceConfigWrapper( + ProtoUtils().compose_proto_from_yaml( + uri=resource_config_uri, + proto_cls=GiglResourceConfig, + ) + ) run_keyword_args = { "job_name": applied_task_identifier, "start_at": start_at, @@ -157,15 +167,12 @@ def run( run_keyword_args["notification_emails"] = notification_emails else: run_keyword_args["notification_emails"] = [ - get_resource_config( - resource_config_uri=resource_config_uri - ).service_account_email + resource_config.service_account_email ] if stop_after is not None: run_keyword_args["stop_after"] = stop_after logger.info(f"Running pipeline with args: {run_keyword_args}") - resource_config = get_resource_config(resource_config_uri=resource_config_uri) vertex_ai_service = VertexAIService( project=resource_config.project, location=resource_config.region, diff --git a/gigl/orchestration/kubeflow/kfp_pipeline.py b/gigl/orchestration/kubeflow/kfp_pipeline.py index b0ad70e30..6fb3f3dce 100644 --- a/gigl/orchestration/kubeflow/kfp_pipeline.py +++ b/gigl/orchestration/kubeflow/kfp_pipeline.py @@ -14,9 +14,6 @@ from gigl.common import LocalUri from gigl.common.logger import Logger from gigl.common.types.resource_config import CommonPipelineComponentConfigs -from gigl.orchestration.kubeflow.utils.glt_backend import ( - check_glt_backend_eligibility_component, -) from gigl.orchestration.kubeflow.utils.log_metrics import log_metrics_to_ui from gigl.orchestration.kubeflow.utils.resource import add_task_resource_requirements from gigl.src.common.constants.components import GiGLComponents @@ -160,16 +157,20 @@ def _generate_component_tasks( resource_config_uri=resource_config_uri, common_pipeline_component_configs=common_pipeline_component_configs, ) - should_use_glt = check_glt_backend_eligibility_component( - task_config_uri=template_or_frozen_config_uri, - base_image=common_pipeline_component_configs.cpu_container_image, - ) + # Composition can depend on sibling files and dynamic resolvers, neither of + # which is represented in KFP's cache key for the primary config URI. + validation_check_task.set_caching_options(enable_caching=False) + resolved_task_config_uri = validation_check_task.outputs["resolved_task_config_uri"] + resolved_resource_config_uri = validation_check_task.outputs[ + "resolved_resource_config_uri" + ] + should_use_glt = validation_check_task.outputs["should_use_glt_backend"] with kfp.dsl.Condition(start_at == GiGLComponents.ConfigPopulator.value): config_populator_task = _create_config_populator_task_op( job_name=job_name, - task_config_uri=template_or_frozen_config_uri, - resource_config_uri=resource_config_uri, + task_config_uri=resolved_task_config_uri, + resource_config_uri=resolved_resource_config_uri, common_pipeline_component_configs=common_pipeline_component_configs, should_use_glt_runtime_param=should_use_glt, stop_after=stop_after, @@ -179,8 +180,8 @@ def _generate_component_tasks( with kfp.dsl.Condition(start_at == GiGLComponents.DataPreprocessor.value): data_preprocessor_task = _create_data_preprocessor_task_op( job_name=job_name, - task_config_uri=template_or_frozen_config_uri, - resource_config_uri=resource_config_uri, + task_config_uri=resolved_task_config_uri, + resource_config_uri=resolved_resource_config_uri, common_pipeline_component_configs=common_pipeline_component_configs, stop_after=stop_after, should_use_glt_runtime_param=should_use_glt, @@ -190,8 +191,8 @@ def _generate_component_tasks( with kfp.dsl.Condition(start_at == GiGLComponents.SubgraphSampler.value): subgraph_sampler_task = _create_subgraph_sampler_task_op( job_name=job_name, - task_config_uri=template_or_frozen_config_uri, - resource_config_uri=resource_config_uri, + task_config_uri=resolved_task_config_uri, + resource_config_uri=resolved_resource_config_uri, common_pipeline_component_configs=common_pipeline_component_configs, stop_after=stop_after, ) @@ -200,8 +201,8 @@ def _generate_component_tasks( with kfp.dsl.Condition(start_at == GiGLComponents.SplitGenerator.value): split_generator_task = _create_split_generator_task_op( job_name=job_name, - task_config_uri=template_or_frozen_config_uri, - resource_config_uri=resource_config_uri, + task_config_uri=resolved_task_config_uri, + resource_config_uri=resolved_resource_config_uri, common_pipeline_component_configs=common_pipeline_component_configs, stop_after=stop_after, ) @@ -210,8 +211,8 @@ def _generate_component_tasks( with kfp.dsl.Condition(start_at == GiGLComponents.Trainer.value): trainer_task = _create_trainer_task_op( job_name=job_name, - task_config_uri=template_or_frozen_config_uri, - resource_config_uri=resource_config_uri, + task_config_uri=resolved_task_config_uri, + resource_config_uri=resolved_resource_config_uri, common_pipeline_component_configs=common_pipeline_component_configs, stop_after=stop_after, ) @@ -220,8 +221,8 @@ def _generate_component_tasks( with kfp.dsl.Condition(start_at == GiGLComponents.Inferencer.value): inferencer_task = _create_inferencer_task_op( job_name=job_name, - task_config_uri=template_or_frozen_config_uri, - resource_config_uri=resource_config_uri, + task_config_uri=resolved_task_config_uri, + resource_config_uri=resolved_resource_config_uri, common_pipeline_component_configs=common_pipeline_component_configs, stop_after=stop_after, ) @@ -230,8 +231,8 @@ def _generate_component_tasks( with kfp.dsl.Condition(start_at == GiGLComponents.PostProcessor.value): post_processor_task = _create_post_processor_task_op( job_name=job_name, - task_config_uri=template_or_frozen_config_uri, - resource_config_uri=resource_config_uri, + task_config_uri=resolved_task_config_uri, + resource_config_uri=resolved_resource_config_uri, common_pipeline_component_configs=common_pipeline_component_configs, ) post_processor_task.after(validation_check_task) diff --git a/gigl/src/validation_check/config_validator.py b/gigl/src/validation_check/config_validator.py index 1e1f77340..0f8c8fdfc 100644 --- a/gigl/src/validation_check/config_validator.py +++ b/gigl/src/validation_check/config_validator.py @@ -1,9 +1,11 @@ import argparse +from pathlib import Path from typing import Optional -from gigl.common import Uri, UriFactory +from gigl.common import GcsUri, Uri, UriFactory from gigl.common.logger import Logger -from gigl.env.pipelines_config import get_resource_config +from gigl.common.utils.gcs import GcsUtils +from gigl.common.utils.proto_utils import ProtoUtils, proto_to_yaml from gigl.src.common.constants.components import GiGLComponents from gigl.src.common.types.pb_wrappers.gbml_config import GbmlConfigPbWrapper from gigl.src.common.types.pb_wrappers.gigl_resource_config import ( @@ -270,31 +272,29 @@ def _run_gbml_and_resource_config_compatibility_checks( ) -def kfp_validation_checks( +def _validate_resolved_configs( job_name: str, - task_config_uri: Uri, start_at: str, - resource_config_uri: Uri, + task_config: gbml_config_pb2.GbmlConfig, + resource_config: GiglResourceConfig, stop_after: Optional[str] = None, -) -> None: +) -> bool: # check if job_name is valid check_if_kfp_pipeline_job_name_valid(job_name=job_name) - # check if start_at and stop_after aligns with live subgraph sampling backend use - check_pipeline_has_valid_start_and_stop_flags( - start_at=start_at, stop_after=stop_after, task_config_uri=task_config_uri.uri - ) - gbml_config_pb_wrapper = GbmlConfigPbWrapper.get_gbml_config_pb_wrapper_from_uri( - gbml_config_uri=task_config_uri - ) - + gbml_config_pb_wrapper = GbmlConfigPbWrapper(task_config) gbml_config_pb: gbml_config_pb2.GbmlConfig = gbml_config_pb_wrapper.gbml_config_pb should_use_live_sgs_backend = gbml_config_pb_wrapper.should_use_glt_backend - resource_config_wrapper: GiglResourceConfigWrapper = get_resource_config( - resource_config_uri=resource_config_uri - ) + resource_config_wrapper = GiglResourceConfigWrapper(resource_config) resource_config_pb: GiglResourceConfig = resource_config_wrapper.resource_config + + # check if start_at and stop_after aligns with live subgraph sampling backend use + check_pipeline_has_valid_start_and_stop_flags( + start_at=start_at, + stop_after=stop_after, + gbml_config_wrapper=gbml_config_pb_wrapper, + ) # check user defined classes and their runtime args if ( @@ -353,6 +353,92 @@ def kfp_validation_checks( assert_trained_model_exists(gbml_config_pb=gbml_config_pb) logger.info("[✅ SUCCESS] All checks passed successfully.") + return should_use_live_sgs_backend + + +def resolve_configs( + task_config_uri: Uri, + resource_config_uri: Uri, +) -> tuple[gbml_config_pb2.GbmlConfig, GiglResourceConfig]: + """Resolve task and resource configs into self-contained protobufs.""" + proto_utils = ProtoUtils() + task_config = proto_utils.compose_proto_from_yaml( + uri=task_config_uri, + proto_cls=gbml_config_pb2.GbmlConfig, + ) + resource_config = proto_utils.compose_proto_from_yaml( + uri=resource_config_uri, + proto_cls=GiglResourceConfig, + ) + if resource_config.WhichOneof("shared_resource") == "shared_resource_config_uri": + # Inline legacy external shared resources so the published snapshot is + # self-contained and downstream components never reread the source URI. + resource_config.shared_resource_config.CopyFrom( + GiglResourceConfigWrapper(resource_config).shared_resource_config + ) + return task_config, resource_config + + +def kfp_validation_checks( + job_name: str, + task_config_uri: Uri, + start_at: str, + resource_config_uri: Uri, + stop_after: Optional[str] = None, +) -> tuple[gbml_config_pb2.GbmlConfig, GiglResourceConfig, bool]: + task_config, resource_config = resolve_configs( + task_config_uri=task_config_uri, + resource_config_uri=resource_config_uri, + ) + should_use_live_sgs_backend = _validate_resolved_configs( + job_name=job_name, + start_at=start_at, + task_config=task_config, + resource_config=resource_config, + stop_after=stop_after, + ) + return task_config, resource_config, should_use_live_sgs_backend + + +def materialize_resolved_configs( + job_name: str, + task_config: gbml_config_pb2.GbmlConfig, + resource_config: GiglResourceConfig, +) -> tuple[GcsUri, GcsUri]: + """Write resolved task and resource configs to stable GCS paths. + + Args: + job_name: Unique name for the pipeline run. + task_config: Resolved task config protobuf. + resource_config: Resolved, self-contained resource config protobuf. + + Returns: + The resolved task and resource config URIs. + """ + resource_config_wrapper = GiglResourceConfigWrapper(resource_config) + snapshot_root = ( + resource_config_wrapper.temp_assets_regional_bucket_path + / job_name + / "config_validator" + ) + task_config_uri = snapshot_root / "resolved_task_config.yaml" + resource_config_uri = snapshot_root / "resolved_resource_config.yaml" + gcs_utils = GcsUtils(project=resource_config_wrapper.project) + gcs_utils.upload_from_string( + gcs_path=task_config_uri, + content=proto_to_yaml(task_config), + ) + gcs_utils.upload_from_string( + gcs_path=resource_config_uri, + content=proto_to_yaml(resource_config), + ) + return task_config_uri, resource_config_uri + + +def _write_kfp_output(path: str, value: str) -> None: + output_path = Path(path) + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_text(value) if __name__ == "__main__": @@ -384,6 +470,24 @@ def kfp_validation_checks( type=str, help="Runtime argument for resource and env specifications of each component", ) + parser.add_argument( + "--output_file_path_resolved_task_config_uri", + type=str, + required=True, + help="KFP output path for the resolved task config URI", + ) + parser.add_argument( + "--output_file_path_resolved_resource_config_uri", + type=str, + required=True, + help="KFP output path for the resolved resource config URI", + ) + parser.add_argument( + "--output_file_path_should_use_glt_backend", + type=str, + required=True, + help="KFP output path for the GLT backend decision", + ) parser.add_argument( "--cpu_docker_uri", type=str, @@ -401,20 +505,48 @@ def kfp_validation_checks( task_config_uri = UriFactory.create_uri(args.task_config_uri) resource_config_uri = UriFactory.create_uri(args.resource_config_uri) - initialize_gigl_runtime( - applied_task_identifier=args.job_name, + check_if_kfp_pipeline_job_name_valid(job_name=args.job_name) + task_config, resource_config = resolve_configs( task_config_uri=task_config_uri, resource_config_uri=resource_config_uri, + ) + ( + resolved_task_config_uri, + resolved_resource_config_uri, + ) = materialize_resolved_configs( + job_name=args.job_name, + task_config=task_config, + resource_config=resource_config, + ) + + # Validation imports user-defined classes. Initialize the historical runtime + # contract first, but point it at the authoritative resolved snapshots. + initialize_gigl_runtime( + applied_task_identifier=args.job_name, + task_config_uri=resolved_task_config_uri, + resource_config_uri=resolved_resource_config_uri, service_name=args.job_name, component=GiGLComponents.ConfigValidator, cpu_docker_uri=args.cpu_docker_uri, cuda_docker_uri=args.cuda_docker_uri, ) - - kfp_validation_checks( + should_use_glt_backend = _validate_resolved_configs( job_name=args.job_name, - task_config_uri=task_config_uri, start_at=args.start_at, - resource_config_uri=resource_config_uri, + task_config=task_config, + resource_config=resource_config, stop_after=args.stop_after, ) + + _write_kfp_output( + args.output_file_path_resolved_task_config_uri, + resolved_task_config_uri.uri, + ) + _write_kfp_output( + args.output_file_path_resolved_resource_config_uri, + resolved_resource_config_uri.uri, + ) + _write_kfp_output( + args.output_file_path_should_use_glt_backend, + "true" if should_use_glt_backend else "false", + ) diff --git a/gigl/src/validation_check/libs/template_config_checks.py b/gigl/src/validation_check/libs/template_config_checks.py index 8b3b86ef6..5bb803d97 100644 --- a/gigl/src/validation_check/libs/template_config_checks.py +++ b/gigl/src/validation_check/libs/template_config_checks.py @@ -42,14 +42,11 @@ def check_if_kfp_pipeline_job_name_valid(job_name: str) -> None: def check_pipeline_has_valid_start_and_stop_flags( start_at: str, stop_after: Optional[str], - task_config_uri: str, + gbml_config_wrapper: GbmlConfigPbWrapper, ) -> None: """ Check if start_at and stop_after are valid with current static (gigl) or dynamic (glt) backend """ - gbml_config_wrapper = GbmlConfigPbWrapper.get_gbml_config_pb_wrapper_from_uri( - gbml_config_uri=UriFactory.create_uri(task_config_uri) - ) components = [start_at] if stop_after is None else [start_at, stop_after] for component in components: if gbml_config_wrapper.should_use_glt_backend: diff --git a/pyproject.toml b/pyproject.toml index 34cbe1129..e26ae17fe 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,6 +25,7 @@ dependencies = [ # See https://stackoverflow.com/questions/69759351/error-jupyter-client-kernelspec-nosuchkernel-no-such-kernel-named-python3-occu "ipykernel", "ipython", + "hydra-core==1.3.2", "kfp>=2.0.0", "matplotlib", "mmh3", @@ -85,9 +86,7 @@ pyg27-torch28-cu128 = [ "torchrec ; sys_platform != 'darwin'", ] -experimental = [ - "hydra-core==1.3.2", -] +experimental = [] # Declares what deps are conflicting so they are not resolved together. # In the case below, we cannot have both pyg27-torch28-cpu and pyg27-torch28-cu128 diff --git a/tests/config_tests/assert_yaml_configs_parse.py b/tests/config_tests/assert_yaml_configs_parse.py index a9958ddf0..2873555d8 100644 --- a/tests/config_tests/assert_yaml_configs_parse.py +++ b/tests/config_tests/assert_yaml_configs_parse.py @@ -31,7 +31,7 @@ import os import re -from gigl.common import Uri, UriFactory +from gigl.common import LocalUri, Uri from gigl.common.logger import Logger from gigl.common.utils.proto_utils import ProtoUtils from snapchat.research.gbml.gbml_config_pb2 import GbmlConfig @@ -80,14 +80,14 @@ def assert_configs_parse(directories: list[str], ignore_regex: list[str] = []) - ) continue total += 1 - yaml_file = UriFactory.create_uri(file_path) + yaml_file = LocalUri(file_path) try: if "resource_config" in file: - proto_utils.read_proto_from_yaml( + proto_utils.compose_proto_from_yaml( yaml_file, GiglResourceConfig ) elif "task_config" in file: - proto_utils.read_proto_from_yaml(yaml_file, GbmlConfig) + proto_utils.compose_proto_from_yaml(yaml_file, GbmlConfig) else: continue logger.info(f"{yaml_file} parsed successfully.") diff --git a/tests/config_tests/notebooks_test.py b/tests/config_tests/notebooks_test.py index d3118b10e..db3075914 100644 --- a/tests/config_tests/notebooks_test.py +++ b/tests/config_tests/notebooks_test.py @@ -8,12 +8,17 @@ import nbformat from nbconvert.preprocessors import ExecutePreprocessor +from gigl.common import LocalUri from gigl.common.constants import GIGL_ROOT_DIR from gigl.common.logger import Logger from gigl.common.types.uri.gcs_uri import GcsUri from gigl.common.types.uri.uri_factory import UriFactory +from gigl.common.utils.proto_utils import ProtoUtils from gigl.env.pipelines_config import get_resource_config -from gigl.src.common.utils.file_loader import FileLoader +from gigl.src.common.types.pb_wrappers.gigl_resource_config import ( + GiglResourceConfigWrapper, +) +from snapchat.research.gbml.gigl_resource_config_pb2 import GiglResourceConfig logger = Logger() @@ -50,19 +55,39 @@ def setUp(self): str(GIGL_ROOT_DIR / "deployment/configs/e2e_glt_resource_config.yaml"), ) logger.info(f"Using resource config URI: {resource_config_uri}") - # Copy over resource config to GCS so all machines can access it. - # If we don't do this, then since different machines run this test case, - # And run some of the buisness logic, they will not have access to the resource config. + # Materialize the resource config in GCS so every machine running + # notebook business logic can access the same self-contained snapshot. + resource_config_source_uri = UriFactory.create_uri(resource_config_uri) + if isinstance( + resource_config_source_uri, LocalUri + ) and resource_config_source_uri.get_basename().endswith(".yaml"): + resource_config = GiglResourceConfigWrapper( + ProtoUtils().compose_proto_from_yaml( + uri=resource_config_source_uri, + proto_cls=GiglResourceConfig, + ) + ) + else: + resource_config = get_resource_config( + resource_config_uri=resource_config_source_uri + ) gcs_uri = GcsUri.join( - get_resource_config().temp_assets_regional_bucket_path, + resource_config.temp_assets_regional_bucket_path, "testing", "notebooks", str(uuid4().hex), "resource_config.yaml", ) logger.info(f"Using GCS URI: {gcs_uri}") - fileloader = FileLoader() - fileloader.load_file(UriFactory.create_uri(resource_config_uri), gcs_uri) + # Upload a self-contained snapshot rather than only the primary source + # file, which may depend on sibling Hydra config fragments. + resource_config.resource_config.shared_resource_config.CopyFrom( + resource_config.shared_resource_config + ) + ProtoUtils().write_proto_to_yaml( + proto=resource_config.resource_config, + uri=gcs_uri, + ) self._notebooks = [ _NoteBookTestConfig( name="cora", diff --git a/tests/unit/common/utils/proto_utils_test.py b/tests/unit/common/utils/proto_utils_test.py index 793e0e7a6..126e5b29a 100644 --- a/tests/unit/common/utils/proto_utils_test.py +++ b/tests/unit/common/utils/proto_utils_test.py @@ -1,13 +1,24 @@ import os +from concurrent.futures import ThreadPoolExecutor from datetime import datetime, timedelta -from tempfile import NamedTemporaryFile +from pathlib import Path +from tempfile import NamedTemporaryFile, TemporaryDirectory +from unittest.mock import patch from absl.testing import absltest +from hydra import initialize_config_dir +from hydra.errors import MissingConfigException +from omegaconf import OmegaConf from gigl.common import LocalUri from gigl.common.logger import Logger +from gigl.common.utils.hydra_config import compose_yaml_config from gigl.common.utils.proto_utils import ProtoUtils from snapchat.research.gbml import gbml_config_pb2 +from snapchat.research.gbml.gigl_resource_config_pb2 import ( + GiglResourceConfig, + SharedResourceConfig, +) from tests.test_assets.test_case import TestCase logger = Logger() @@ -82,6 +93,258 @@ def test_read_proto_from_yaml_raises_typeerror_when_root_is_not_a_mapping(self): finally: os.remove(tmp_file.name) + def test_can_compose_task_config_from_primary_parent(self): + with TemporaryDirectory() as temp_directory: + config_root = Path(temp_directory) + (config_root / "shared").mkdir() + (config_root / "task.yaml").write_text( + "defaults:\n" + " - shared@sharedConfig: directed\n" + " - _self_\n" + "datasetConfig:\n" + " dataPreprocessorConfig:\n" + " dataPreprocessorArgs:\n" + ' yesterday: "${now:%Y%m%d,days-1}"\n' + ) + (config_root / "shared" / "directed.yaml").write_text( + "isGraphDirected: true\n" + ) + + task_config = self.proto_utils.compose_proto_from_yaml( + uri=LocalUri(config_root / "task.yaml"), + proto_cls=gbml_config_pb2.GbmlConfig, + ) + + self.assertTrue(task_config.shared_config.is_graph_directed) + self.assertEqual( + task_config.dataset_config.data_preprocessor_config.data_preprocessor_args[ + "yesterday" + ], + (datetime.now() - timedelta(days=1)).strftime("%Y%m%d"), + ) + + def test_can_compose_resource_config_from_primary_parent(self): + with TemporaryDirectory() as temp_directory: + config_root = Path(temp_directory) + (config_root / "shared").mkdir() + (config_root / "resource.yaml").write_text( + "defaults:\n - shared@shared_resource_config: local\n - _self_\n" + ) + (config_root / "shared" / "local.yaml").write_text( + "common_compute_config:\n" + " project: example-project\n" + " region: us-central1\n" + " temp_regional_assets_bucket: gs://example-bucket\n" + ) + + resource_config = self.proto_utils.compose_proto_from_yaml( + uri=LocalUri(config_root / "resource.yaml"), + proto_cls=GiglResourceConfig, + ) + + self.assertEqual( + resource_config.shared_resource_config.common_compute_config.project, + "example-project", + ) + + def test_can_compose_any_proto(self): + with TemporaryDirectory() as temp_directory: + config_root = Path(temp_directory) + (config_root / "compute").mkdir() + config_path = config_root / "shared_resource.yaml" + config_path.write_text( + "defaults:\n - compute@common_compute_config: local\n - _self_\n" + ) + (config_root / "compute" / "local.yaml").write_text( + "project: example-project\nregion: us-central1\n" + ) + + shared_resource_config = self.proto_utils.compose_proto_from_yaml( + uri=LocalUri(config_path), + proto_cls=SharedResourceConfig, + ) + + self.assertEqual( + shared_resource_config.common_compute_config.project, + "example-project", + ) + + def test_resource_config_resolves_dynamic_interpolation(self): + with TemporaryDirectory() as temp_directory: + config_path = Path(temp_directory) / "resource.yaml" + config_path.write_text( + "shared_resource_config:\n" + " common_compute_config:\n" + ' project: "${oc.env:PROJECT_ID}"\n' + ) + + with patch.dict(os.environ, {"PROJECT_ID": "example-project"}): + resource_config = self.proto_utils.read_proto_from_yaml( + uri=LocalUri(config_path), + proto_cls=GiglResourceConfig, + ) + + self.assertEqual( + resource_config.shared_resource_config.common_compute_config.project, + "example-project", + ) + + def test_resource_config_allows_deterministic_value_interpolation(self): + with TemporaryDirectory() as temp_directory: + config_path = Path(temp_directory) / "resource.yaml" + config_path.write_text( + "shared_resource_config:\n" + " common_compute_config:\n" + " project: example-project\n" + ' region: "${.project}"\n' + ) + + resource_config = self.proto_utils.read_proto_from_yaml( + uri=LocalUri(config_path), + proto_cls=GiglResourceConfig, + ) + + self.assertEqual( + resource_config.shared_resource_config.common_compute_config.region, + "example-project", + ) + + def test_resource_composition_ignores_dynamic_resolver_in_unselected_file(self): + with TemporaryDirectory() as temp_directory: + config_root = Path(temp_directory) + (config_root / "shared").mkdir() + config_path = config_root / "resource.yaml" + config_path.write_text( + "defaults:\n - shared@shared_resource_config: local\n - _self_\n" + ) + (config_root / "shared" / "local.yaml").write_text( + "common_compute_config:\n" + " project: example-project\n" + " region: us-central1\n" + ) + (config_root / "unselected_task.yaml").write_text( + 'run_name: "${now:%Y%m%d}"\n' + ) + + resource_config = self.proto_utils.compose_proto_from_yaml( + uri=LocalUri(config_path), + proto_cls=GiglResourceConfig, + ) + + self.assertEqual( + resource_config.shared_resource_config.common_compute_config.project, + "example-project", + ) + + def test_resource_composition_resolves_dynamic_resolver_in_selected_file(self): + with TemporaryDirectory() as temp_directory: + config_root = Path(temp_directory) + (config_root / "shared").mkdir() + config_path = config_root / "resource.yaml" + config_path.write_text( + "defaults:\n - shared@shared_resource_config: dynamic\n - _self_\n" + ) + (config_root / "shared" / "dynamic.yaml").write_text( + 'common_compute_config:\n project: "${oc.env:PROJECT_ID}"\n' + ) + + with patch.dict(os.environ, {"PROJECT_ID": "example-project"}): + resource_config = self.proto_utils.compose_proto_from_yaml( + uri=LocalUri(config_path), + proto_cls=GiglResourceConfig, + ) + + self.assertEqual( + resource_config.shared_resource_config.common_compute_config.project, + "example-project", + ) + + def test_resource_composition_resolves_dynamic_nested_default(self): + with TemporaryDirectory() as temp_directory: + config_root = Path(temp_directory) + (config_root / "base").mkdir() + (config_root / "compute").mkdir() + config_path = config_root / "resource.yaml" + config_path.write_text( + "defaults:\n - base@_global_: resource\n - _self_\n" + ) + (config_root / "base" / "resource.yaml").write_text( + "defaults:\n" + " - /compute@shared_resource_config: " + "${oc.env:RESOURCE_PROFILE}\n" + " - _self_\n" + ) + (config_root / "compute" / "local.yaml").write_text( + "common_compute_config:\n" + " project: example-project\n" + " region: us-central1\n" + ) + + with patch.dict(os.environ, {"RESOURCE_PROFILE": "local"}): + resource_config = self.proto_utils.compose_proto_from_yaml( + uri=LocalUri(config_path), + proto_cls=GiglResourceConfig, + ) + + self.assertEqual( + resource_config.shared_resource_config.common_compute_config.project, + "example-project", + ) + + def test_composition_is_thread_safe(self): + with TemporaryDirectory() as temp_directory: + config_root = Path(temp_directory) + (config_root / "shared").mkdir() + config_path = config_root / "task.yaml" + config_path.write_text( + "defaults:\n - shared@sharedConfig: directed\n - _self_\n" + ) + (config_root / "shared" / "directed.yaml").write_text( + "isGraphDirected: true\n" + ) + + def read_config(_: int) -> bool: + return ( + ProtoUtils() + .compose_proto_from_yaml( + uri=LocalUri(config_path), + proto_cls=gbml_config_pb2.GbmlConfig, + ) + .shared_config.is_graph_directed + ) + + with ThreadPoolExecutor(max_workers=8) as executor: + results = list(executor.map(read_config, range(32))) + + self.assertEqual(results, [True] * 32) + + def test_composition_preserves_foreign_hydra_context(self): + with TemporaryDirectory() as temp_directory: + config_root = Path(temp_directory) + config_path = config_root / "task.yaml" + config_path.write_text("defaults:\n - _self_\n") + + with initialize_config_dir( + config_dir=str(config_root), + version_base="1.3", + ): + with self.assertRaises(RuntimeError): + self.proto_utils.compose_proto_from_yaml( + uri=LocalUri(config_path), + proto_cls=gbml_config_pb2.GbmlConfig, + ) + + def test_composition_restores_gigl_resolvers_after_failure(self): + with TemporaryDirectory() as temp_directory: + config_path = Path(temp_directory) / "task.yaml" + config_path.write_text("defaults:\n - missing\n") + + with self.assertRaises(MissingConfigException): + compose_yaml_config(LocalUri(config_path)) + + config = OmegaConf.create({"tomorrow": "${now:%Y-%m-%d,days+1}"}) + self.assertRegex(config.tomorrow, r"\d{4}-\d{2}-\d{2}") + if __name__ == "__main__": absltest.main() diff --git a/tests/unit/orchestration/kubeflow/component_spec_test.py b/tests/unit/orchestration/kubeflow/component_spec_test.py index 6e705b7d4..922f9c550 100644 --- a/tests/unit/orchestration/kubeflow/component_spec_test.py +++ b/tests/unit/orchestration/kubeflow/component_spec_test.py @@ -13,6 +13,25 @@ def _read_component_spec(component_name: str) -> str: class ComponentSpecTest(TestCase): + def test_config_validator_resolves_both_configs(self) -> None: + component_spec = _read_component_spec("config_validator") + + self.assertIn("name: resolved_task_config_uri", component_spec) + self.assertIn("name: resolved_resource_config_uri", component_spec) + self.assertIn("name: should_use_glt_backend", component_spec) + self.assertIn( + "{outputPath: resolved_task_config_uri}", + component_spec, + ) + self.assertIn( + "{outputPath: resolved_resource_config_uri}", + component_spec, + ) + self.assertIn( + "{outputPath: should_use_glt_backend}", + component_spec, + ) + def test_data_preprocessor_accepts_source_image_uri_flags(self) -> None: component_spec = _read_component_spec("data_preprocessor") diff --git a/tests/unit/orchestration/kubeflow/kfp_pipeline_test.py b/tests/unit/orchestration/kubeflow/kfp_pipeline_test.py new file mode 100644 index 000000000..5d4652816 --- /dev/null +++ b/tests/unit/orchestration/kubeflow/kfp_pipeline_test.py @@ -0,0 +1,50 @@ +from pathlib import Path +from tempfile import TemporaryDirectory + +import yaml +from absl.testing import absltest +from kfp.compiler import Compiler + +from gigl.common.types.resource_config import CommonPipelineComponentConfigs +from gigl.orchestration.kubeflow.kfp_pipeline import generate_pipeline +from tests.test_assets.test_case import TestCase + + +class KfpPipelineTest(TestCase): + def test_validator_outputs_are_pipeline_inputs_for_downstream_tasks(self) -> None: + pipeline = generate_pipeline( + CommonPipelineComponentConfigs( + cuda_container_image="cuda-image", + cpu_container_image="cpu-image", + dataflow_container_image="dataflow-image", + ) + ) + with TemporaryDirectory() as temp_directory: + output_path = Path(temp_directory) / "pipeline.yaml" + Compiler().compile(pipeline, str(output_path)) + pipeline_text = output_path.read_text() + pipeline_spec = yaml.safe_load(pipeline_text) + + validator_task = pipeline_spec["components"]["comp-exit-handler-1"]["dag"][ + "tasks" + ]["kfp-validation-check"] + self.assertEqual(validator_task["cachingOptions"], {}) + self.assertNotIn("check-glt-backend", pipeline_text) + self.assertIn( + "outputParameterKey: resolved_task_config_uri\n" + " producerTask: kfp-validation-check", + pipeline_text, + ) + self.assertIn( + "outputParameterKey: resolved_resource_config_uri\n" + " producerTask: kfp-validation-check", + pipeline_text, + ) + self.assertIn( + "outputParameterKey: should_use_glt_backend", + pipeline_text, + ) + + +if __name__ == "__main__": + absltest.main() diff --git a/tests/unit/src/validation/config_validator_test.py b/tests/unit/src/validation/config_validator_test.py index 5373803f7..71e7d5058 100644 --- a/tests/unit/src/validation/config_validator_test.py +++ b/tests/unit/src/validation/config_validator_test.py @@ -1,15 +1,21 @@ import os import shutil import tempfile +from pathlib import Path +from unittest.mock import patch import google.protobuf.message from absl.testing import absltest from parameterized import param, parameterized import gigl.env.pipelines_config -from gigl.common import Uri, UriFactory +from gigl.common import GcsUri, Uri, UriFactory from gigl.common.utils.proto_utils import ProtoUtils -from gigl.src.validation_check.config_validator import kfp_validation_checks +from gigl.src.validation_check.config_validator import ( + kfp_validation_checks, + materialize_resolved_configs, + resolve_configs, +) from snapchat.research.gbml import ( gbml_config_pb2, gigl_resource_config_pb2, @@ -351,6 +357,130 @@ def test_resource_config_validation_failure_with_mock_configs( resource_config_uri=self._live_resource_config_uri, ) + @patch("gigl.src.validation_check.config_validator.GcsUtils") + def test_composes_and_materializes_both_configs(self, mock_gcs_utils) -> None: + config_root = Path(self._temp_dir) / "composed" + (config_root / "task").mkdir(parents=True) + (config_root / "resource").mkdir() + + task_config_path = config_root / "task_config.yaml" + task_config_path.write_text("defaults:\n - task@_global_: base\n - _self_\n") + (config_root / "task" / "base.yaml").write_text( + "graph_metadata:\n" + " node_types: [paper]\n" + "shared_config:\n" + " is_graph_directed: true\n" + ) + + shared_resource_config_path = config_root / "shared_resource_config.yaml" + shared_resource_config_path.write_text( + "common_compute_config:\n" + " project: ${oc.env:CONFIG_VALIDATOR_TEST_PROJECT}\n" + " region: us-central1\n" + " temp_regional_assets_bucket: gs://test-temp-regional\n" + ) + resource_config_path = config_root / "resource_config.yaml" + resource_config_path.write_text( + "defaults:\n - resource@_global_: base\n - _self_\n" + ) + (config_root / "resource" / "base.yaml").write_text( + f"shared_resource_config_uri: {shared_resource_config_path}\n" + ) + + with patch.dict( + os.environ, + {"CONFIG_VALIDATOR_TEST_PROJECT": "composed-project"}, + ): + task_config, resource_config = resolve_configs( + task_config_uri=UriFactory.create_uri(str(task_config_path)), + resource_config_uri=UriFactory.create_uri(str(resource_config_path)), + ) + + self.assertTrue(task_config.shared_config.is_graph_directed) + self.assertEqual( + resource_config.WhichOneof("shared_resource"), + "shared_resource_config", + ) + self.assertEqual( + resource_config.shared_resource_config.common_compute_config.project, + "composed-project", + ) + + snapshot_directory = Path(self._temp_dir) / "snapshots" + snapshot_directory.mkdir() + + def write_snapshot(gcs_path, content): + (snapshot_directory / gcs_path.get_basename()).write_text(content) + + mock_gcs_utils.return_value.upload_from_string.side_effect = write_snapshot + task_uri, resource_uri = materialize_resolved_configs( + job_name="config-resolution-test", + task_config=task_config, + resource_config=resource_config, + ) + + self.assertEqual(task_uri.get_basename(), "resolved_task_config.yaml") + self.assertEqual(resource_uri.get_basename(), "resolved_resource_config.yaml") + materialized_task_config = self._proto_utils.read_proto_from_yaml( + UriFactory.create_uri( + str(snapshot_directory / "resolved_task_config.yaml") + ), + gbml_config_pb2.GbmlConfig, + ) + materialized_resource_config = self._proto_utils.read_proto_from_yaml( + UriFactory.create_uri( + str(snapshot_directory / "resolved_resource_config.yaml") + ), + gigl_resource_config_pb2.GiglResourceConfig, + ) + self.assertEqual(materialized_task_config, task_config) + self.assertEqual(materialized_resource_config, resource_config) + + @patch("gigl.common.utils.proto_utils.FileLoader.load_file") + def test_downloads_remote_configs_before_composing(self, mock_load_file) -> None: + task_uri = GcsUri("gs://test-configs/task.yaml") + resource_uri = GcsUri("gs://test-configs/resource.yaml") + search_root = Path(self._temp_dir) / "remote_shared_configs" + (search_root / "task").mkdir(parents=True) + (search_root / "task" / "base.yaml").write_text( + "shared_config:\n is_graph_directed: true\n" + ) + source_configs = { + task_uri.uri: ( + "hydra:\n" + " searchpath:\n" + f" - file://{search_root}\n" + "defaults:\n" + " - task@_global_: base\n" + " - _self_\n" + ), + resource_uri.uri: ( + "shared_resource_config:\n" + " common_compute_config:\n" + " project: remote-project\n" + " temp_regional_assets_bucket: gs://test-temp-regional\n" + ), + } + + def download_config( + file_uri_src, + file_uri_dst, + ): + self.assertTrue(file_uri_dst.get_basename().endswith(".yaml")) + Path(file_uri_dst.uri).write_text(source_configs[file_uri_src.uri]) + + mock_load_file.side_effect = download_config + task_config, resource_config = resolve_configs( + task_config_uri=task_uri, + resource_config_uri=resource_uri, + ) + + self.assertTrue(task_config.shared_config.is_graph_directed) + self.assertEqual( + resource_config.shared_resource_config.common_compute_config.project, + "remote-project", + ) + if __name__ == "__main__": absltest.main() diff --git a/uv.lock b/uv.lock index e0c67d828..598bd5b49 100644 --- a/uv.lock +++ b/uv.lock @@ -713,6 +713,7 @@ dependencies = [ { name = "google-cloud-logging" }, { name = "google-cloud-pipeline-components" }, { name = "google-cloud-storage" }, + { name = "hydra-core" }, { name = "ipykernel" }, { name = "ipython" }, { name = "kfp" }, @@ -733,9 +734,6 @@ dependencies = [ ] [package.optional-dependencies] -experimental = [ - { name = "hydra-core" }, -] pyg27-torch28-cpu = [ { name = "fbgemm-gpu", version = "1.3.0+cpu", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "sys_platform != 'darwin'" }, { name = "pyg-lib", version = "0.5.0+pt28cpu", source = { registry = "https://data.pyg.org/whl/torch-2.8.0+cpu.html" }, marker = "sys_platform != 'darwin'" }, @@ -852,7 +850,7 @@ requires-dist = [ { name = "google-cloud-logging" }, { name = "google-cloud-pipeline-components" }, { name = "google-cloud-storage" }, - { name = "hydra-core", marker = "extra == 'experimental'", specifier = "==1.3.2" }, + { name = "hydra-core", specifier = "==1.3.2" }, { name = "ipykernel" }, { name = "ipython" }, { name = "kfp", specifier = ">=2.0.0" },