From ae6fd44d92e0f3edb747f480136b1d7bf01ae06b Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Thu, 30 Jul 2026 16:22:47 +0000 Subject: [PATCH 1/9] Add Hydra composition for GiGL configs --- deployment/configs/e2e/preprocessor.yaml | 11 + deployment/configs/e2e/shared.yaml | 15 + .../configs/e2e_cicd_resource_config.yaml | 37 +- .../configs/e2e_glt_gs_resource_config.yaml | 30 +- .../configs/e2e_glt_resource_config.yaml | 30 +- .../config_guides/hydra_composition.md | 97 ++++++ .../config_guides/resource_config_guide.md | 3 + .../config_guides/task_config_guide.md | 3 + gigl/common/omegaconf_resolvers.py | 13 +- gigl/common/utils/hydra_config.py | 216 ++++++++++++ gigl/common/utils/proto_utils.py | 61 +++- .../config_validator/component.yaml | 10 +- .../kubeflow/kfp_orchestrator.py | 43 ++- gigl/orchestration/kubeflow/kfp_pipeline.py | 51 +-- .../types/pb_wrappers/gigl_resource_config.py | 16 + gigl/src/validation_check/config_validator.py | 204 ++++++++++-- .../libs/template_config_checks.py | 5 +- pyproject.toml | 5 +- tests/config_tests/notebooks_test.py | 20 +- tests/unit/common/utils/proto_utils_test.py | 315 +++++++++++++++++- .../kubeflow/component_spec_test.py | 20 ++ .../kubeflow/kfp_pipeline_test.py | 54 +++ .../src/validation/config_validator_test.py | 64 +++- uv.lock | 6 +- 24 files changed, 1167 insertions(+), 162 deletions(-) create mode 100644 deployment/configs/e2e/preprocessor.yaml create mode 100644 deployment/configs/e2e/shared.yaml create mode 100644 docs/user_guide/config_guides/hydra_composition.md create mode 100644 gigl/common/utils/hydra_config.py create mode 100644 tests/unit/orchestration/kubeflow/kfp_pipeline_test.py 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..0f01b9d7b --- /dev/null +++ b/docs/user_guide/config_guides/hydra_composition.md @@ -0,0 +1,97 @@ +# Composing task and resource configs with Hydra + +GiGL task and resource configs can opt into [Hydra Defaults Lists](https://hydra.cc/docs/advanced/defaults_list/) by +adding a top-level `defaults` list. Configs without `defaults` continue through the existing YAML and OmegaConf loader. + +## Config root + +The primary file's parent directory is its Hydra config root. Composed primary files and selected fragments must use +Hydra's `.yaml` extension. Plain legacy `.yml` files without `defaults` remain supported. + +```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`. Composition currently requires a local primary file. GCS and HTTP primaries remain supported only as plain, +single-file YAML without a Defaults List. + +## 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 +``` + +Resource config is read once before pipeline submission to select the project, region, service account, and staging +bucket. ConfigValidator resolves it again inside the pipeline and verifies that both resolutions have the same protobuf +fingerprint. Dynamic resolvers such as `now`, `git_hash`, and `oc.env` are therefore rejected in resource configs. +Deterministic references to another config value remain supported. + +## 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. +- Defaults selectors must be static; interpolation is supported only in composed values. +- Parent traversal and external `hydra.searchpath` entries are rejected. +- GCS and HTTP configs remain supported only when they are plain, single-file YAML. +- 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/omegaconf_resolvers.py b/gigl/common/omegaconf_resolvers.py index 6653620e5..841d061c1 100644 --- a/gigl/common/omegaconf_resolvers.py +++ b/gigl/common/omegaconf_resolvers.py @@ -157,23 +157,26 @@ def git_hash_resolver() -> str: return "" -def register_resolvers() -> None: +def register_resolvers(replace: bool = False) -> None: """Register all custom OmegaConf resolvers. This function should be called at application startup to register all custom resolvers with OmegaConf. + + Args: + replace: Whether to replace resolvers registered by another library. """ - if not OmegaConf.has_resolver("now"): + if replace or not OmegaConf.has_resolver("now"): logger.info("Registering OmegaConf resolver 'now'") - OmegaConf.register_new_resolver("now", now_resolver) + OmegaConf.register_new_resolver("now", now_resolver, replace=replace) else: logger.debug( "OmegaConf resolver 'now' already registered, skipping registration" ) - if not OmegaConf.has_resolver("git_hash"): + if replace or not OmegaConf.has_resolver("git_hash"): logger.info("Registering OmegaConf resolver 'git_hash'") - OmegaConf.register_new_resolver("git_hash", git_hash_resolver) + OmegaConf.register_new_resolver("git_hash", git_hash_resolver, replace=replace) else: logger.debug( "OmegaConf resolver 'git_hash' already registered, skipping registration" diff --git a/gigl/common/utils/hydra_config.py b/gigl/common/utils/hydra_config.py new file mode 100644 index 000000000..604b8b6c6 --- /dev/null +++ b/gigl/common/utils/hydra_config.py @@ -0,0 +1,216 @@ +"""Hydra composition support for GiGL protobuf YAML configs.""" + +from __future__ import annotations + +import os +import re +import threading +from collections.abc import Mapping, Sequence +from pathlib import Path, PurePosixPath +from typing import Any, cast + +import yaml +from hydra import compose, initialize_config_dir +from hydra.core.global_hydra import GlobalHydra +from omegaconf import OmegaConf +from omegaconf.basecontainer import BaseContainer + +from gigl.common import HttpUri, LocalUri, Uri +from gigl.common.omegaconf_resolvers import register_resolvers + +_COMPOSE_LOCK = threading.RLock() +_MAX_CONFIG_FILE_COUNT = 1_000 +_MAX_CONFIG_TOTAL_BYTES = 50 * 1024 * 1024 +_YAML_SUFFIXES = (".yaml", ".yml") +_RESOLVER_INTERPOLATION = re.compile(r"\$\{[^{}]+:") +_HYDRA_MUTATED_RESOLVERS = ( + "now", + "hydra", + "python_version", + "git_hash", + "oc.env", +) + + +def contains_dynamic_interpolation(value: Any) -> bool: + """Return whether YAML contains a resolver that can vary by process.""" + if isinstance(value, str): + return bool(_RESOLVER_INTERPOLATION.search(value)) or "${git_hash}" in value + if isinstance(value, Mapping): + return any( + contains_dynamic_interpolation(key) or contains_dynamic_interpolation(item) + for key, item in value.items() + ) + if isinstance(value, Sequence) and not isinstance(value, (str, bytes)): + return any(contains_dynamic_interpolation(item) for item in value) + return False + + +def compose_yaml_config( + uri: Uri, + *, + reject_interpolations: bool = False, +) -> dict[str, Any]: + """Compose a YAML config with Hydra using its parent as the config root. + + Args: + uri: Primary YAML config URI. + reject_interpolations: Whether to reject OmegaConf interpolations before + Hydra resolves them. + + Returns: + A fully composed and resolved mapping. + + Raises: + TypeError: If composition is requested for an unsupported URI. + ValueError: If the bundle is unsafe or resolves to a non-mapping value. + RuntimeError: If another Hydra application owns the global context. + """ + primary_name = uri.get_basename() + if not primary_name.endswith(".yaml"): + raise ValueError(f"Hydra primary config must use the .yaml extension: {uri}") + config_name = primary_name.rsplit(".", 1)[0] + + if isinstance(uri, LocalUri): + config_root = Path(uri.uri).absolute().parent + _validate_local_bundle(config_root=config_root) + return _compose_local( + config_root=config_root, + config_name=config_name, + reject_interpolations=reject_interpolations, + ) + + if isinstance(uri, HttpUri): + raise TypeError("Hydra composition is not supported for HTTP config URIs.") + raise TypeError(f"Hydra composition is not supported for {type(uri).__name__}.") + + +def _compose_local( + config_root: Path, + config_name: str, + reject_interpolations: bool, +) -> dict[str, Any]: + with _COMPOSE_LOCK: + if GlobalHydra.instance().is_initialized(): + raise RuntimeError( + "GiGL cannot compose a config while another Hydra context is active." + ) + # Hydra mutates OmegaConf's process-global resolver registry and does not + # restore it when its context exits. OmegaConf exposes no public way to + # recover the prior resolver callables, so preserve the affected wrappers + # directly and restore them in ``finally``. + prior_resolvers = { + name: BaseContainer._resolvers.get(name) + for name in _HYDRA_MUTATED_RESOLVERS + } + 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(replace=True) + if reject_interpolations: + for resolver_name in ("now", "git_hash", "oc.env"): + if resolver_name in BaseContainer._resolvers: + BaseContainer._resolvers[resolver_name] = ( + _reject_resource_resolver + ) + composed = compose(config_name=config_name, overrides=[]) + if reject_interpolations and contains_dynamic_interpolation( + OmegaConf.to_container(composed, resolve=False) + ): + raise ValueError( + "Resource configs cannot contain dynamic OmegaConf resolvers " + "because submission and pipeline validation run in separate " + "processes." + ) + resolved = OmegaConf.to_container(composed, resolve=True) + finally: + for name, resolver in prior_resolvers.items(): + if resolver is None: + BaseContainer._resolvers.pop(name, None) + else: + BaseContainer._resolvers[name] = resolver + + 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) + + +def _reject_resource_resolver(*_: Any) -> Any: + raise ValueError( + "Resource configs cannot use process-dependent OmegaConf resolvers." + ) + + +def _validate_local_bundle(config_root: Path) -> None: + resolved_root = config_root.resolve() + yaml_file_count = 0 + total_bytes = 0 + for yaml_path in config_root.rglob("*"): + if yaml_path.is_symlink(): + raise ValueError(f"Config bundle {config_root} contains a symlink.") + if not yaml_path.is_file() or yaml_path.suffix not in _YAML_SUFFIXES: + continue + yaml_file_count += 1 + if yaml_file_count > _MAX_CONFIG_FILE_COUNT: + raise ValueError( + f"Config bundle {config_root} contains more than " + f"{_MAX_CONFIG_FILE_COUNT} YAML files." + ) + total_bytes += yaml_path.stat().st_size + if total_bytes > _MAX_CONFIG_TOTAL_BYTES: + raise ValueError( + f"Config bundle {config_root} exceeds the " + f"{_MAX_CONFIG_TOTAL_BYTES}-byte limit." + ) + if not yaml_path.resolve().is_relative_to(resolved_root): + raise ValueError(f"Config bundle {config_root} contains an unsafe path.") + with yaml_path.open("r") as file: + raw_data = yaml.safe_load(file) + _validate_hydra_controls(raw_data=raw_data, source_path=yaml_path) + + +def _validate_hydra_controls(raw_data: Any, source_path: Path) -> None: + if not isinstance(raw_data, Mapping): + return + hydra_config = raw_data.get("hydra") + if isinstance(hydra_config, Mapping) and "searchpath" in hydra_config: + raise ValueError( + f"External Hydra search paths are not supported in {source_path}." + ) + defaults = raw_data.get("defaults") + if not isinstance(defaults, Sequence) or isinstance(defaults, (str, bytes)): + return + for default in defaults: + selectors: list[str] = [] + if isinstance(default, str): + selectors.append(default) + elif isinstance(default, Mapping): + selectors.extend(str(key) for key in default) + selectors.extend( + str(value) for value in default.values() if isinstance(value, str) + ) + for selector in selectors: + if "${" in selector: + raise ValueError( + f"Interpolations in Hydra defaults are not supported in " + f"{source_path}." + ) + selector_path = selector.split("@", 1)[0] + selector_parts = selector_path.split() + while selector_parts and selector_parts[0] in {"optional", "override"}: + selector_parts.pop(0) + selector_path = selector_parts[-1] if selector_parts else "" + if ".." in PurePosixPath(selector_path).parts: + raise ValueError( + f"Parent traversal in Hydra defaults is not supported in " + f"{source_path}." + ) diff --git a/gigl/common/utils/proto_utils.py b/gigl/common/utils/proto_utils.py index e61258756..c42fb5f3c 100644 --- a/gigl/common/utils/proto_utils.py +++ b/gigl/common/utils/proto_utils.py @@ -1,3 +1,4 @@ +import hashlib from tempfile import NamedTemporaryFile from typing import Optional, Type, TypeVar, cast @@ -9,12 +10,50 @@ 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, + contains_dynamic_interpolation, +) from gigl.src.common.utils.file_loader import FileLoader logger = Logger() T = TypeVar("T", bound=message.Message) +_HYDRA_PROTO_TYPES = { + "snapchat.research.gbml.GbmlConfig", + "snapchat.research.gbml.GiglResourceConfig", +} +_DETERMINISTIC_RESOURCE_PROTO_TYPES = { + "snapchat.research.gbml.GiglResourceConfig", + "snapchat.research.gbml.SharedResourceConfig", +} + + +def get_proto_fingerprint(proto: message.Message) -> str: + """Return a stable SHA-256 fingerprint for a protobuf message. + + Args: + proto: Protobuf message to fingerprint. + + Returns: + Hexadecimal SHA-256 digest of the deterministic protobuf bytes. + """ + return hashlib.sha256(proto.SerializeToString(deterministic=True)).hexdigest() + + +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: @@ -25,9 +64,27 @@ def read_proto_from_yaml(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, "r") as file: raw_data = yaml.safe_load(file) - omega_conf_obj = OmegaConf.create(raw_data) tfh.close() - obj_dict = OmegaConf.to_object(omega_conf_obj) + proto_type = proto_cls.DESCRIPTOR.full_name + reject_interpolations = proto_type in _DETERMINISTIC_RESOURCE_PROTO_TYPES + if reject_interpolations and contains_dynamic_interpolation(raw_data): + raise ValueError( + "Resource configs cannot contain dynamic OmegaConf resolvers because " + "submission and pipeline validation run in separate processes." + ) + + if ( + isinstance(raw_data, dict) + and "defaults" in raw_data + and proto_type in _HYDRA_PROTO_TYPES + ): + obj_dict = compose_yaml_config( + uri=uri, + reject_interpolations=reject_interpolations, + ) + else: + omega_conf_obj = OmegaConf.create(raw_data) + obj_dict = OmegaConf.to_object(omega_conf_obj) if not isinstance(obj_dict, dict): raise TypeError( f"ProtoUtils.read_proto_from_yaml expected a mapping at the YAML root for " diff --git a/gigl/orchestration/kubeflow/components/config_validator/component.yaml b/gigl/orchestration/kubeflow/components/config_validator/component.yaml index 7c4ec057e..31302ff8f 100644 --- a/gigl/orchestration/kubeflow/components/config_validator/component.yaml +++ b/gigl/orchestration/kubeflow/components/config_validator/component.yaml @@ -5,10 +5,14 @@ inputs: - {name: task_config_uri, type: String, description: 'Template or frozen GBML config uri'} - {name: start_at, type: String, description: 'Start component'} - {name: resource_config_uri, type: String, description: 'Runtine argument for resource and env specifications of each component'} +- {name: bootstrap_resource_config_hash, type: String, description: 'Resource config fingerprint computed during pipeline submission'} - {name: stop_after, type: String, description: 'Stop component'} - {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: @@ -19,7 +23,11 @@ implementation: --task_config_uri, {inputValue: task_config_uri}, --start_at, {inputValue: start_at}, --resource_config_uri, {inputValue: resource_config_uri}, + --bootstrap_resource_config_hash, {inputValue: bootstrap_resource_config_hash}, --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..ebbdcc9f0 100644 --- a/gigl/orchestration/kubeflow/kfp_orchestrator.py +++ b/gigl/orchestration/kubeflow/kfp_orchestrator.py @@ -4,6 +4,7 @@ from pathlib import Path from typing import Optional, Union +import yaml from google.cloud import aiplatform from kfp.compiler import Compiler @@ -12,15 +13,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, get_proto_fingerprint 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() @@ -36,6 +41,7 @@ ) DEFAULT_START_AT_COMPONENT = "config_populator" +_CONFIG_RESOLUTION_PIPELINE_PARAMETER = "bootstrap_resource_config_hash" class KfpOrchestrator: @@ -141,13 +147,43 @@ def run( assert file_loader.does_uri_exist(compiled_pipeline_path), ( f"Compiled pipeline path {compiled_pipeline_path} does not exist." ) + compiled_pipeline_file = file_loader.load_to_temp_file( + file_uri_src=compiled_pipeline_path, + delete=True, + should_create_symlinks_if_possible=False, + ) + with open(compiled_pipeline_file.name, "r") as file: + compiled_pipeline = yaml.safe_load(file) + compiled_pipeline_file.close() + pipeline_parameters = ( + compiled_pipeline.get("root", {}) + .get("inputDefinitions", {}) + .get("parameters", {}) + if isinstance(compiled_pipeline, dict) + else {} + ) + if _CONFIG_RESOLUTION_PIPELINE_PARAMETER not in pipeline_parameters: + raise ValueError( + "The compiled pipeline predates config resolution support. " + "Recompile the pipeline before running it." + ) logger.info(f"Skipping pipeline compilation; will use {compiled_pipeline_path}") + resource_config = GiglResourceConfigWrapper( + ProtoUtils().read_proto_from_yaml( + uri=resource_config_uri, + proto_cls=GiglResourceConfig, + ) + ) + resolved_resource_config = resource_config.get_resolved_resource_config() run_keyword_args = { "job_name": applied_task_identifier, "start_at": start_at, "template_or_frozen_config_uri": task_config_uri.uri, "resource_config_uri": resource_config_uri.uri, + "bootstrap_resource_config_hash": get_proto_fingerprint( + resolved_resource_config + ), } # We need to provide *some* notification emails, other wise the cleanup component will fail. # Ideally, we'd be able to provide None and have it handle it, but for whatever reason @@ -157,15 +193,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..3f42a8c48 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 @@ -59,6 +56,7 @@ def _generate_component_task( common_pipeline_component_configs: CommonPipelineComponentConfigs, start_at: Optional[str] = None, stop_after: Optional[str] = None, + bootstrap_resource_config_hash: Optional[str] = None, ) -> PipelineTask: component_task: PipelineTask if component == GiGLComponents.ConfigPopulator: @@ -72,11 +70,16 @@ def _generate_component_task( ) elif component == GiGLComponents.ConfigValidator: + if bootstrap_resource_config_hash is None: + raise ValueError( + "ConfigValidator requires a bootstrap resource config hash." + ) component_task = _speced_component_op_dict[component]( job_name=job_name, task_config_uri=task_config_uri, start_at=start_at, resource_config_uri=resource_config_uri, + bootstrap_resource_config_hash=bootstrap_resource_config_hash, stop_after=stop_after, cpu_docker_uri=common_pipeline_component_configs.cpu_container_image, cuda_docker_uri=common_pipeline_component_configs.cuda_container_image, @@ -147,6 +150,7 @@ def _generate_component_tasks( job_name: str, template_or_frozen_config_uri: str, resource_config_uri: str, + bootstrap_resource_config_hash: str, common_pipeline_component_configs: CommonPipelineComponentConfigs, start_at: Optional[str] = None, stop_after: Optional[str] = None, @@ -158,18 +162,21 @@ def _generate_component_tasks( start_at=start_at, stop_after=stop_after, resource_config_uri=resource_config_uri, + bootstrap_resource_config_hash=bootstrap_resource_config_hash, 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, - ) + 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 +186,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 +197,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 +207,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 +217,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 +227,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 +237,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) @@ -260,6 +267,7 @@ def pipeline( job_name: str, template_or_frozen_config_uri: str, resource_config_uri: str, + bootstrap_resource_config_hash: str, start_at: str = GiGLComponents.ConfigPopulator.value, stop_after: Optional[str] = None, notification_emails: Optional[List[str]] = None, @@ -272,6 +280,7 @@ def pipeline( job_name=job_name, template_or_frozen_config_uri=template_or_frozen_config_uri, resource_config_uri=resource_config_uri, + bootstrap_resource_config_hash=bootstrap_resource_config_hash, common_pipeline_component_configs=common_pipeline_component_configs, start_at=start_at, stop_after=stop_after, diff --git a/gigl/src/common/types/pb_wrappers/gigl_resource_config.py b/gigl/src/common/types/pb_wrappers/gigl_resource_config.py index 6a5d0b40b..a2809c927 100644 --- a/gigl/src/common/types/pb_wrappers/gigl_resource_config.py +++ b/gigl/src/common/types/pb_wrappers/gigl_resource_config.py @@ -97,6 +97,22 @@ def shared_resource_config(self) -> SharedResourceConfig: return self._loaded_shared_resource_config + def get_resolved_resource_config(self) -> GiglResourceConfig: + """Return a self-contained resource config. + + If the source config references a shared resource config URI, the + referenced protobuf is copied into the returned message. + + Returns: + A new resource config with an inline shared resource config. + """ + resolved_resource_config = GiglResourceConfig() + resolved_resource_config.CopyFrom(self.resource_config) + resolved_resource_config.shared_resource_config.CopyFrom( + self.shared_resource_config + ) + return resolved_resource_config + @property def get_resource_config_uri(self) -> str: parser = argparse.ArgumentParser() diff --git a/gigl/src/validation_check/config_validator.py b/gigl/src/validation_check/config_validator.py index 1e1f77340..092c065d6 100644 --- a/gigl/src/validation_check/config_validator.py +++ b/gigl/src/validation_check/config_validator.py @@ -1,9 +1,15 @@ 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, + get_proto_fingerprint, + 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 +276,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 +357,107 @@ 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, + bootstrap_resource_config_hash: Optional[str] = None, +) -> tuple[gbml_config_pb2.GbmlConfig, GiglResourceConfig]: + """Resolve task and resource configs into self-contained protobufs.""" + proto_utils = ProtoUtils() + task_config = proto_utils.read_proto_from_yaml( + uri=task_config_uri, + proto_cls=gbml_config_pb2.GbmlConfig, + ) + source_resource_config_wrapper = GiglResourceConfigWrapper( + proto_utils.read_proto_from_yaml( + uri=resource_config_uri, + proto_cls=GiglResourceConfig, + ) + ) + resource_config = source_resource_config_wrapper.get_resolved_resource_config() + + resolved_resource_config_hash = get_proto_fingerprint(resource_config) + if ( + bootstrap_resource_config_hash is not None + and bootstrap_resource_config_hash != resolved_resource_config_hash + ): + raise ValueError( + "Resource config changed between pipeline submission and validation." + ) + 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, + bootstrap_resource_config_hash: 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, + bootstrap_resource_config_hash=bootstrap_resource_config_hash, + ) + 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 = GcsUri.join( + resource_config_wrapper.temp_assets_regional_bucket_path, + job_name, + "config_validator", + ) + task_config_uri = GcsUri.join( + snapshot_root, + f"resolved_task_config_{get_proto_fingerprint(task_config)}.yaml", + ) + resource_config_uri = GcsUri.join( + snapshot_root, + f"resolved_resource_config_{get_proto_fingerprint(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 +489,30 @@ def kfp_validation_checks( type=str, help="Runtime argument for resource and env specifications of each component", ) + parser.add_argument( + "--bootstrap_resource_config_hash", + type=str, + required=True, + help="Resource config fingerprint computed during pipeline submission", + ) + 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 +530,49 @@ 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, + bootstrap_resource_config_hash=args.bootstrap_resource_config_hash, + ) + ( + 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/notebooks_test.py b/tests/config_tests/notebooks_test.py index d3118b10e..7dd1df456 100644 --- a/tests/config_tests/notebooks_test.py +++ b/tests/config_tests/notebooks_test.py @@ -12,8 +12,8 @@ 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 logger = Logger() @@ -50,19 +50,25 @@ 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 = get_resource_config( + resource_config_uri=UriFactory.create_uri(resource_config_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. + ProtoUtils().write_proto_to_yaml( + proto=resource_config.get_resolved_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..7a20c9411 100644 --- a/tests/unit/common/utils/proto_utils_test.py +++ b/tests/unit/common/utils/proto_utils_test.py @@ -1,13 +1,21 @@ 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 omegaconf import OmegaConf -from gigl.common import LocalUri +from gigl.common import GcsUri, LocalUri from gigl.common.logger import Logger +from gigl.common.omegaconf_resolvers import now_resolver +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 from tests.test_assets.test_case import TestCase logger = Logger() @@ -82,6 +90,309 @@ 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.read_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.read_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_resource_config_rejects_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 self.assertRaises(ValueError): + self.proto_utils.read_proto_from_yaml( + uri=LocalUri(config_path), + proto_cls=GiglResourceConfig, + ) + + 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.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_composition_rejects_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 self.assertRaises(ValueError): + self.proto_utils.read_proto_from_yaml( + uri=LocalUri(config_path), + proto_cls=GiglResourceConfig, + ) + + def test_resource_composition_rejects_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"}), + self.assertRaises(ValueError), + ): + self.proto_utils.read_proto_from_yaml( + uri=LocalUri(config_path), + proto_cls=GiglResourceConfig, + ) + + def test_composed_primary_requires_yaml_extension(self): + with TemporaryDirectory() as temp_directory: + config_path = Path(temp_directory) / "task.yml" + config_path.write_text("defaults:\n - _self_\n") + + with self.assertRaises(ValueError): + self.proto_utils.read_proto_from_yaml( + uri=LocalUri(config_path), + proto_cls=gbml_config_pb2.GbmlConfig, + ) + + def test_gcs_composition_is_not_supported(self): + with self.assertRaisesRegex( + TypeError, + "Hydra composition is not supported for GcsUri", + ): + compose_yaml_config(GcsUri("gs://example-bucket/configs/task.yaml")) + + 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() + .read_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.read_proto_from_yaml( + uri=LocalUri(config_path), + proto_cls=gbml_config_pb2.GbmlConfig, + ) + + def test_composition_restores_foreign_omegaconf_resolver(self): + with TemporaryDirectory() as temp_directory: + config_path = Path(temp_directory) / "task.yaml" + config_path.write_text("defaults:\n - _self_\n") + OmegaConf.register_new_resolver( + "now", lambda *_: "foreign-now", replace=True + ) + try: + self.proto_utils.read_proto_from_yaml( + uri=LocalUri(config_path), + proto_cls=gbml_config_pb2.GbmlConfig, + ) + + config = OmegaConf.create({"value": "${now:any-format}"}) + self.assertEqual(config.value, "foreign-now") + finally: + OmegaConf.register_new_resolver("now", now_resolver, replace=True) + + def test_composition_rejects_parent_traversal(self): + for selector in ( + "../outside@_global_", + "optional ../outside@_global_", + "override optional ../outside@_global_", + ): + with self.subTest(selector=selector): + with TemporaryDirectory() as temp_directory: + config_root = Path(temp_directory) / "bundle" + config_root.mkdir() + config_path = config_root / "task.yaml" + config_path.write_text(f"defaults:\n - {selector}\n") + (Path(temp_directory) / "outside.yaml").write_text( + "sharedConfig:\n isGraphDirected: true\n" + ) + + with self.assertRaises(ValueError): + self.proto_utils.read_proto_from_yaml( + uri=LocalUri(config_path), + proto_cls=gbml_config_pb2.GbmlConfig, + ) + + def test_composition_rejects_interpolated_default_selector(self): + with TemporaryDirectory() as temp_directory: + config_path = Path(temp_directory) / "task.yaml" + config_path.write_text( + "defaults:\n - group@sharedConfig: ${oc.env:TASK_PROFILE}\n" + ) + + with ( + patch.dict(os.environ, {"TASK_PROFILE": "../outside"}), + self.assertRaises(ValueError), + ): + self.proto_utils.read_proto_from_yaml( + uri=LocalUri(config_path), + proto_cls=gbml_config_pb2.GbmlConfig, + ) + + def test_composition_rejects_symlinked_config_group(self): + with TemporaryDirectory() as temp_directory: + temp_root = Path(temp_directory) + config_root = temp_root / "bundle" + outside_group = temp_root / "outside" + config_root.mkdir() + outside_group.mkdir() + config_path = config_root / "task.yaml" + config_path.write_text( + "defaults:\n - escaped@sharedConfig: directed\n - _self_\n" + ) + (outside_group / "directed.yaml").write_text("isGraphDirected: true\n") + (config_root / "escaped").symlink_to( + outside_group, target_is_directory=True + ) + + with self.assertRaises(ValueError): + self.proto_utils.read_proto_from_yaml( + uri=LocalUri(config_path), + proto_cls=gbml_config_pb2.GbmlConfig, + ) + 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..b5dec6670 100644 --- a/tests/unit/orchestration/kubeflow/component_spec_test.py +++ b/tests/unit/orchestration/kubeflow/component_spec_test.py @@ -13,6 +13,26 @@ 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: bootstrap_resource_config_hash", component_spec) + 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..9d2757793 --- /dev/null +++ b/tests/unit/orchestration/kubeflow/kfp_pipeline_test.py @@ -0,0 +1,54 @@ +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) + + self.assertIn( + "bootstrap_resource_config_hash", + pipeline_spec["root"]["inputDefinitions"]["parameters"], + ) + 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..f95750384 100644 --- a/tests/unit/src/validation/config_validator_test.py +++ b/tests/unit/src/validation/config_validator_test.py @@ -1,6 +1,7 @@ import os import shutil import tempfile +from unittest.mock import patch import google.protobuf.message from absl.testing import absltest @@ -8,8 +9,11 @@ import gigl.env.pipelines_config from gigl.common import Uri, UriFactory -from gigl.common.utils.proto_utils import ProtoUtils -from gigl.src.validation_check.config_validator import kfp_validation_checks +from gigl.common.utils.proto_utils import ProtoUtils, get_proto_fingerprint +from gigl.src.validation_check.config_validator import ( + kfp_validation_checks, + materialize_resolved_configs, +) from snapchat.research.gbml import ( gbml_config_pb2, gigl_resource_config_pb2, @@ -351,6 +355,62 @@ def test_resource_config_validation_failure_with_mock_configs( resource_config_uri=self._live_resource_config_uri, ) + def test_resource_config_bootstrap_hash_must_match(self) -> None: + with self.assertRaises(ValueError): + kfp_validation_checks( + job_name="resource_config_validation_test", + task_config_uri=self._live_task_config_uri, + start_at="config_populator", + resource_config_uri=self._live_resource_config_uri, + bootstrap_resource_config_hash="stale-resource-config", + ) + + def test_resource_config_bootstrap_hash_accepts_resolved_config(self) -> None: + resource_config = self._proto_utils.read_proto_from_yaml( + uri=self._live_resource_config_uri, + proto_cls=gigl_resource_config_pb2.GiglResourceConfig, + ) + + _, resolved_resource_config, should_use_glt_backend = kfp_validation_checks( + job_name="resource_config_validation_test", + task_config_uri=self._live_task_config_uri, + start_at="config_populator", + resource_config_uri=self._live_resource_config_uri, + bootstrap_resource_config_hash=get_proto_fingerprint(resource_config), + ) + + self.assertTrue(should_use_glt_backend) + self.assertEqual( + get_proto_fingerprint(resolved_resource_config), + get_proto_fingerprint(resource_config), + ) + + @patch("gigl.src.validation_check.config_validator.GcsUtils") + def test_materializes_both_resolved_configs(self, mock_gcs_utils) -> None: + task_config = _create_valid_live_subgraph_sampling_task_config() + resource_config = _create_valid_live_subgraph_sampling_resource_config() + + task_uri, resource_uri = materialize_resolved_configs( + job_name="config-resolution-test", + task_config=task_config, + resource_config=resource_config, + ) + + self.assertTrue( + task_uri.uri.startswith( + "gs://test-temp-regional/config-resolution-test/config_validator/" + ) + ) + self.assertTrue( + resource_uri.uri.startswith( + "gs://test-temp-regional/config-resolution-test/config_validator/" + ) + ) + self.assertEqual( + mock_gcs_utils.return_value.upload_from_string.call_count, + 2, + ) + 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" }, From 3ebd26d2b5a56f69beaacdd7740cd41c4bda0489 Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Thu, 30 Jul 2026 16:35:36 +0000 Subject: [PATCH 2/9] Simplify Hydra composition wrapper --- .../config_guides/hydra_composition.md | 2 - gigl/common/utils/hydra_config.py | 125 +++--------------- gigl/common/utils/proto_utils.py | 8 +- tests/unit/common/utils/proto_utils_test.py | 63 +-------- 4 files changed, 29 insertions(+), 169 deletions(-) diff --git a/docs/user_guide/config_guides/hydra_composition.md b/docs/user_guide/config_guides/hydra_composition.md index 0f01b9d7b..05da459cc 100644 --- a/docs/user_guide/config_guides/hydra_composition.md +++ b/docs/user_guide/config_guides/hydra_composition.md @@ -91,7 +91,5 @@ and runner intentionally differ. ## Boundaries - GiGL does not consume Hydra command-line overrides, multirun, launchers, or output-directory behavior. -- Defaults selectors must be static; interpolation is supported only in composed values. -- Parent traversal and external `hydra.searchpath` entries are rejected. - GCS and HTTP configs remain supported only when they are plain, single-file YAML. - Treat config bundle write access as trusted access. GiGL configs can reference importable classes and commands. diff --git a/gigl/common/utils/hydra_config.py b/gigl/common/utils/hydra_config.py index 604b8b6c6..b1c288dba 100644 --- a/gigl/common/utils/hydra_config.py +++ b/gigl/common/utils/hydra_config.py @@ -6,10 +6,9 @@ import re import threading from collections.abc import Mapping, Sequence -from pathlib import Path, PurePosixPath +from pathlib import Path from typing import Any, cast -import yaml from hydra import compose, initialize_config_dir from hydra.core.global_hydra import GlobalHydra from omegaconf import OmegaConf @@ -19,23 +18,14 @@ from gigl.common.omegaconf_resolvers import register_resolvers _COMPOSE_LOCK = threading.RLock() -_MAX_CONFIG_FILE_COUNT = 1_000 -_MAX_CONFIG_TOTAL_BYTES = 50 * 1024 * 1024 -_YAML_SUFFIXES = (".yaml", ".yml") -_RESOLVER_INTERPOLATION = re.compile(r"\$\{[^{}]+:") -_HYDRA_MUTATED_RESOLVERS = ( - "now", - "hydra", - "python_version", - "git_hash", - "oc.env", -) +_DYNAMIC_RESOLVER_INTERPOLATION = re.compile(r"\$\{(?:now|git_hash|oc\.env)(?::|})") +_DYNAMIC_RESOLVERS = ("now", "git_hash", "oc.env") def contains_dynamic_interpolation(value: Any) -> bool: """Return whether YAML contains a resolver that can vary by process.""" if isinstance(value, str): - return bool(_RESOLVER_INTERPOLATION.search(value)) or "${git_hash}" in value + return bool(_DYNAMIC_RESOLVER_INTERPOLATION.search(value)) if isinstance(value, Mapping): return any( contains_dynamic_interpolation(key) or contains_dynamic_interpolation(item) @@ -49,21 +39,22 @@ def contains_dynamic_interpolation(value: Any) -> bool: def compose_yaml_config( uri: Uri, *, - reject_interpolations: bool = False, + reject_dynamic_interpolations: bool = False, ) -> dict[str, Any]: """Compose a YAML config with Hydra using its parent as the config root. Args: uri: Primary YAML config URI. - reject_interpolations: Whether to reject OmegaConf interpolations before - Hydra resolves them. + reject_dynamic_interpolations: Whether to reject process-dependent + OmegaConf resolvers before Hydra resolves them. Returns: A fully composed and resolved mapping. Raises: TypeError: If composition is requested for an unsupported URI. - ValueError: If the bundle is unsafe or resolves to a non-mapping value. + ValueError: If the primary filename is unsupported or the result is not + a mapping. RuntimeError: If another Hydra application owns the global context. """ primary_name = uri.get_basename() @@ -73,11 +64,10 @@ def compose_yaml_config( if isinstance(uri, LocalUri): config_root = Path(uri.uri).absolute().parent - _validate_local_bundle(config_root=config_root) return _compose_local( config_root=config_root, config_name=config_name, - reject_interpolations=reject_interpolations, + reject_dynamic_interpolations=reject_dynamic_interpolations, ) if isinstance(uri, HttpUri): @@ -88,7 +78,7 @@ def compose_yaml_config( def _compose_local( config_root: Path, config_name: str, - reject_interpolations: bool, + reject_dynamic_interpolations: bool, ) -> dict[str, Any]: with _COMPOSE_LOCK: if GlobalHydra.instance().is_initialized(): @@ -99,10 +89,7 @@ def _compose_local( # restore it when its context exits. OmegaConf exposes no public way to # recover the prior resolver callables, so preserve the affected wrappers # directly and restore them in ``finally``. - prior_resolvers = { - name: BaseContainer._resolvers.get(name) - for name in _HYDRA_MUTATED_RESOLVERS - } + prior_resolvers = BaseContainer._resolvers.copy() try: with initialize_config_dir( config_dir=os.fspath(config_root), @@ -113,14 +100,15 @@ def _compose_local( # initialization. GiGL's resolver is a backward-compatible # superset that also supports offsets. register_resolvers(replace=True) - if reject_interpolations: - for resolver_name in ("now", "git_hash", "oc.env"): - if resolver_name in BaseContainer._resolvers: - BaseContainer._resolvers[resolver_name] = ( - _reject_resource_resolver - ) + if reject_dynamic_interpolations: + for resolver_name in _DYNAMIC_RESOLVERS: + OmegaConf.register_new_resolver( + resolver_name, + _reject_resource_resolver, + replace=True, + ) composed = compose(config_name=config_name, overrides=[]) - if reject_interpolations and contains_dynamic_interpolation( + if reject_dynamic_interpolations and contains_dynamic_interpolation( OmegaConf.to_container(composed, resolve=False) ): raise ValueError( @@ -130,11 +118,8 @@ def _compose_local( ) resolved = OmegaConf.to_container(composed, resolve=True) finally: - for name, resolver in prior_resolvers.items(): - if resolver is None: - BaseContainer._resolvers.pop(name, None) - else: - BaseContainer._resolvers[name] = resolver + BaseContainer._resolvers.clear() + BaseContainer._resolvers.update(prior_resolvers) if not isinstance(resolved, dict): raise ValueError( @@ -148,69 +133,3 @@ def _reject_resource_resolver(*_: Any) -> Any: raise ValueError( "Resource configs cannot use process-dependent OmegaConf resolvers." ) - - -def _validate_local_bundle(config_root: Path) -> None: - resolved_root = config_root.resolve() - yaml_file_count = 0 - total_bytes = 0 - for yaml_path in config_root.rglob("*"): - if yaml_path.is_symlink(): - raise ValueError(f"Config bundle {config_root} contains a symlink.") - if not yaml_path.is_file() or yaml_path.suffix not in _YAML_SUFFIXES: - continue - yaml_file_count += 1 - if yaml_file_count > _MAX_CONFIG_FILE_COUNT: - raise ValueError( - f"Config bundle {config_root} contains more than " - f"{_MAX_CONFIG_FILE_COUNT} YAML files." - ) - total_bytes += yaml_path.stat().st_size - if total_bytes > _MAX_CONFIG_TOTAL_BYTES: - raise ValueError( - f"Config bundle {config_root} exceeds the " - f"{_MAX_CONFIG_TOTAL_BYTES}-byte limit." - ) - if not yaml_path.resolve().is_relative_to(resolved_root): - raise ValueError(f"Config bundle {config_root} contains an unsafe path.") - with yaml_path.open("r") as file: - raw_data = yaml.safe_load(file) - _validate_hydra_controls(raw_data=raw_data, source_path=yaml_path) - - -def _validate_hydra_controls(raw_data: Any, source_path: Path) -> None: - if not isinstance(raw_data, Mapping): - return - hydra_config = raw_data.get("hydra") - if isinstance(hydra_config, Mapping) and "searchpath" in hydra_config: - raise ValueError( - f"External Hydra search paths are not supported in {source_path}." - ) - defaults = raw_data.get("defaults") - if not isinstance(defaults, Sequence) or isinstance(defaults, (str, bytes)): - return - for default in defaults: - selectors: list[str] = [] - if isinstance(default, str): - selectors.append(default) - elif isinstance(default, Mapping): - selectors.extend(str(key) for key in default) - selectors.extend( - str(value) for value in default.values() if isinstance(value, str) - ) - for selector in selectors: - if "${" in selector: - raise ValueError( - f"Interpolations in Hydra defaults are not supported in " - f"{source_path}." - ) - selector_path = selector.split("@", 1)[0] - selector_parts = selector_path.split() - while selector_parts and selector_parts[0] in {"optional", "override"}: - selector_parts.pop(0) - selector_path = selector_parts[-1] if selector_parts else "" - if ".." in PurePosixPath(selector_path).parts: - raise ValueError( - f"Parent traversal in Hydra defaults is not supported in " - f"{source_path}." - ) diff --git a/gigl/common/utils/proto_utils.py b/gigl/common/utils/proto_utils.py index c42fb5f3c..2f9e4d1f0 100644 --- a/gigl/common/utils/proto_utils.py +++ b/gigl/common/utils/proto_utils.py @@ -66,8 +66,10 @@ def read_proto_from_yaml(self, uri: Uri, proto_cls: Type[T]) -> T: raw_data = yaml.safe_load(file) tfh.close() proto_type = proto_cls.DESCRIPTOR.full_name - reject_interpolations = proto_type in _DETERMINISTIC_RESOURCE_PROTO_TYPES - if reject_interpolations and contains_dynamic_interpolation(raw_data): + reject_dynamic_interpolations = ( + proto_type in _DETERMINISTIC_RESOURCE_PROTO_TYPES + ) + if reject_dynamic_interpolations and contains_dynamic_interpolation(raw_data): raise ValueError( "Resource configs cannot contain dynamic OmegaConf resolvers because " "submission and pipeline validation run in separate processes." @@ -80,7 +82,7 @@ def read_proto_from_yaml(self, uri: Uri, proto_cls: Type[T]) -> T: ): obj_dict = compose_yaml_config( uri=uri, - reject_interpolations=reject_interpolations, + reject_dynamic_interpolations=reject_dynamic_interpolations, ) else: omega_conf_obj = OmegaConf.create(raw_data) diff --git a/tests/unit/common/utils/proto_utils_test.py b/tests/unit/common/utils/proto_utils_test.py index 7a20c9411..e7d0d852b 100644 --- a/tests/unit/common/utils/proto_utils_test.py +++ b/tests/unit/common/utils/proto_utils_test.py @@ -7,6 +7,7 @@ from absl.testing import absltest from hydra import initialize_config_dir +from hydra.errors import ConfigCompositionException from omegaconf import OmegaConf from gigl.common import GcsUri, LocalUri @@ -247,7 +248,7 @@ def test_resource_composition_rejects_dynamic_nested_default(self): with ( patch.dict(os.environ, {"RESOURCE_PROFILE": "local"}), - self.assertRaises(ValueError), + self.assertRaises(ConfigCompositionException), ): self.proto_utils.read_proto_from_yaml( uri=LocalUri(config_path), @@ -333,66 +334,6 @@ def test_composition_restores_foreign_omegaconf_resolver(self): finally: OmegaConf.register_new_resolver("now", now_resolver, replace=True) - def test_composition_rejects_parent_traversal(self): - for selector in ( - "../outside@_global_", - "optional ../outside@_global_", - "override optional ../outside@_global_", - ): - with self.subTest(selector=selector): - with TemporaryDirectory() as temp_directory: - config_root = Path(temp_directory) / "bundle" - config_root.mkdir() - config_path = config_root / "task.yaml" - config_path.write_text(f"defaults:\n - {selector}\n") - (Path(temp_directory) / "outside.yaml").write_text( - "sharedConfig:\n isGraphDirected: true\n" - ) - - with self.assertRaises(ValueError): - self.proto_utils.read_proto_from_yaml( - uri=LocalUri(config_path), - proto_cls=gbml_config_pb2.GbmlConfig, - ) - - def test_composition_rejects_interpolated_default_selector(self): - with TemporaryDirectory() as temp_directory: - config_path = Path(temp_directory) / "task.yaml" - config_path.write_text( - "defaults:\n - group@sharedConfig: ${oc.env:TASK_PROFILE}\n" - ) - - with ( - patch.dict(os.environ, {"TASK_PROFILE": "../outside"}), - self.assertRaises(ValueError), - ): - self.proto_utils.read_proto_from_yaml( - uri=LocalUri(config_path), - proto_cls=gbml_config_pb2.GbmlConfig, - ) - - def test_composition_rejects_symlinked_config_group(self): - with TemporaryDirectory() as temp_directory: - temp_root = Path(temp_directory) - config_root = temp_root / "bundle" - outside_group = temp_root / "outside" - config_root.mkdir() - outside_group.mkdir() - config_path = config_root / "task.yaml" - config_path.write_text( - "defaults:\n - escaped@sharedConfig: directed\n - _self_\n" - ) - (outside_group / "directed.yaml").write_text("isGraphDirected: true\n") - (config_root / "escaped").symlink_to( - outside_group, target_is_directory=True - ) - - with self.assertRaises(ValueError): - self.proto_utils.read_proto_from_yaml( - uri=LocalUri(config_path), - proto_cls=gbml_config_pb2.GbmlConfig, - ) - if __name__ == "__main__": absltest.main() From 0b7adc7375e5bc8a64609c120d173a3b51b951a9 Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Thu, 30 Jul 2026 16:50:48 +0000 Subject: [PATCH 3/9] Remove defensive config resolution machinery --- .../config_guides/hydra_composition.md | 6 +- gigl/common/utils/hydra_config.py | 75 ++----------------- gigl/common/utils/proto_utils.py | 35 +-------- .../config_validator/component.yaml | 2 - .../kubeflow/kfp_orchestrator.py | 39 +--------- gigl/orchestration/kubeflow/kfp_pipeline.py | 10 --- gigl/src/validation_check/config_validator.py | 29 +------ tests/unit/common/utils/proto_utils_test.py | 59 +++++++-------- .../kubeflow/component_spec_test.py | 1 - .../kubeflow/kfp_pipeline_test.py | 4 - .../src/validation/config_validator_test.py | 48 +++--------- 11 files changed, 49 insertions(+), 259 deletions(-) diff --git a/docs/user_guide/config_guides/hydra_composition.md b/docs/user_guide/config_guides/hydra_composition.md index 05da459cc..d837ae392 100644 --- a/docs/user_guide/config_guides/hydra_composition.md +++ b/docs/user_guide/config_guides/hydra_composition.md @@ -60,9 +60,9 @@ common_compute_config: ``` Resource config is read once before pipeline submission to select the project, region, service account, and staging -bucket. ConfigValidator resolves it again inside the pipeline and verifies that both resolutions have the same protobuf -fingerprint. Dynamic resolvers such as `now`, `git_hash`, and `oc.env` are therefore rejected in resource configs. -Deterministic references to another config value remain supported. +bucket. ConfigValidator resolves 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 read. ## Pipeline behavior diff --git a/gigl/common/utils/hydra_config.py b/gigl/common/utils/hydra_config.py index b1c288dba..03b32c1d7 100644 --- a/gigl/common/utils/hydra_config.py +++ b/gigl/common/utils/hydra_config.py @@ -3,50 +3,25 @@ from __future__ import annotations import os -import re import threading -from collections.abc import Mapping, Sequence 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 omegaconf.basecontainer import BaseContainer -from gigl.common import HttpUri, LocalUri, Uri +from gigl.common import LocalUri, Uri from gigl.common.omegaconf_resolvers import register_resolvers _COMPOSE_LOCK = threading.RLock() -_DYNAMIC_RESOLVER_INTERPOLATION = re.compile(r"\$\{(?:now|git_hash|oc\.env)(?::|})") -_DYNAMIC_RESOLVERS = ("now", "git_hash", "oc.env") -def contains_dynamic_interpolation(value: Any) -> bool: - """Return whether YAML contains a resolver that can vary by process.""" - if isinstance(value, str): - return bool(_DYNAMIC_RESOLVER_INTERPOLATION.search(value)) - if isinstance(value, Mapping): - return any( - contains_dynamic_interpolation(key) or contains_dynamic_interpolation(item) - for key, item in value.items() - ) - if isinstance(value, Sequence) and not isinstance(value, (str, bytes)): - return any(contains_dynamic_interpolation(item) for item in value) - return False - - -def compose_yaml_config( - uri: Uri, - *, - reject_dynamic_interpolations: bool = False, -) -> dict[str, Any]: +def compose_yaml_config(uri: Uri) -> dict[str, Any]: """Compose a YAML config with Hydra using its parent as the config root. Args: uri: Primary YAML config URI. - reject_dynamic_interpolations: Whether to reject process-dependent - OmegaConf resolvers before Hydra resolves them. Returns: A fully composed and resolved mapping. @@ -64,32 +39,14 @@ def compose_yaml_config( if isinstance(uri, LocalUri): config_root = Path(uri.uri).absolute().parent - return _compose_local( - config_root=config_root, - config_name=config_name, - reject_dynamic_interpolations=reject_dynamic_interpolations, - ) + else: + raise TypeError(f"Hydra composition is not supported for {type(uri).__name__}.") - if isinstance(uri, HttpUri): - raise TypeError("Hydra composition is not supported for HTTP config URIs.") - raise TypeError(f"Hydra composition is not supported for {type(uri).__name__}.") - - -def _compose_local( - config_root: Path, - config_name: str, - reject_dynamic_interpolations: bool, -) -> dict[str, Any]: with _COMPOSE_LOCK: if GlobalHydra.instance().is_initialized(): raise RuntimeError( "GiGL cannot compose a config while another Hydra context is active." ) - # Hydra mutates OmegaConf's process-global resolver registry and does not - # restore it when its context exits. OmegaConf exposes no public way to - # recover the prior resolver callables, so preserve the affected wrappers - # directly and restore them in ``finally``. - prior_resolvers = BaseContainer._resolvers.copy() try: with initialize_config_dir( config_dir=os.fspath(config_root), @@ -100,26 +57,10 @@ def _compose_local( # initialization. GiGL's resolver is a backward-compatible # superset that also supports offsets. register_resolvers(replace=True) - if reject_dynamic_interpolations: - for resolver_name in _DYNAMIC_RESOLVERS: - OmegaConf.register_new_resolver( - resolver_name, - _reject_resource_resolver, - replace=True, - ) composed = compose(config_name=config_name, overrides=[]) - if reject_dynamic_interpolations and contains_dynamic_interpolation( - OmegaConf.to_container(composed, resolve=False) - ): - raise ValueError( - "Resource configs cannot contain dynamic OmegaConf resolvers " - "because submission and pipeline validation run in separate " - "processes." - ) resolved = OmegaConf.to_container(composed, resolve=True) finally: - BaseContainer._resolvers.clear() - BaseContainer._resolvers.update(prior_resolvers) + register_resolvers(replace=True) if not isinstance(resolved, dict): raise ValueError( @@ -127,9 +68,3 @@ def _compose_local( f"{type(resolved).__name__}, expected a mapping." ) return cast(dict[str, Any], resolved) - - -def _reject_resource_resolver(*_: Any) -> Any: - raise ValueError( - "Resource configs cannot use process-dependent OmegaConf resolvers." - ) diff --git a/gigl/common/utils/proto_utils.py b/gigl/common/utils/proto_utils.py index 2f9e4d1f0..9037b320d 100644 --- a/gigl/common/utils/proto_utils.py +++ b/gigl/common/utils/proto_utils.py @@ -1,4 +1,3 @@ -import hashlib from tempfile import NamedTemporaryFile from typing import Optional, Type, TypeVar, cast @@ -10,10 +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, - contains_dynamic_interpolation, -) +from gigl.common.utils.hydra_config import compose_yaml_config from gigl.src.common.utils.file_loader import FileLoader logger = Logger() @@ -24,22 +20,6 @@ "snapchat.research.gbml.GbmlConfig", "snapchat.research.gbml.GiglResourceConfig", } -_DETERMINISTIC_RESOURCE_PROTO_TYPES = { - "snapchat.research.gbml.GiglResourceConfig", - "snapchat.research.gbml.SharedResourceConfig", -} - - -def get_proto_fingerprint(proto: message.Message) -> str: - """Return a stable SHA-256 fingerprint for a protobuf message. - - Args: - proto: Protobuf message to fingerprint. - - Returns: - Hexadecimal SHA-256 digest of the deterministic protobuf bytes. - """ - return hashlib.sha256(proto.SerializeToString(deterministic=True)).hexdigest() def proto_to_yaml(proto: message.Message) -> str: @@ -66,24 +46,13 @@ def read_proto_from_yaml(self, uri: Uri, proto_cls: Type[T]) -> T: raw_data = yaml.safe_load(file) tfh.close() proto_type = proto_cls.DESCRIPTOR.full_name - reject_dynamic_interpolations = ( - proto_type in _DETERMINISTIC_RESOURCE_PROTO_TYPES - ) - if reject_dynamic_interpolations and contains_dynamic_interpolation(raw_data): - raise ValueError( - "Resource configs cannot contain dynamic OmegaConf resolvers because " - "submission and pipeline validation run in separate processes." - ) if ( isinstance(raw_data, dict) and "defaults" in raw_data and proto_type in _HYDRA_PROTO_TYPES ): - obj_dict = compose_yaml_config( - uri=uri, - reject_dynamic_interpolations=reject_dynamic_interpolations, - ) + obj_dict = compose_yaml_config(uri=uri) else: omega_conf_obj = OmegaConf.create(raw_data) obj_dict = OmegaConf.to_object(omega_conf_obj) diff --git a/gigl/orchestration/kubeflow/components/config_validator/component.yaml b/gigl/orchestration/kubeflow/components/config_validator/component.yaml index 31302ff8f..fc1963c4a 100644 --- a/gigl/orchestration/kubeflow/components/config_validator/component.yaml +++ b/gigl/orchestration/kubeflow/components/config_validator/component.yaml @@ -5,7 +5,6 @@ inputs: - {name: task_config_uri, type: String, description: 'Template or frozen GBML config uri'} - {name: start_at, type: String, description: 'Start component'} - {name: resource_config_uri, type: String, description: 'Runtine argument for resource and env specifications of each component'} -- {name: bootstrap_resource_config_hash, type: String, description: 'Resource config fingerprint computed during pipeline submission'} - {name: stop_after, type: String, description: 'Stop component'} - {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"} @@ -23,7 +22,6 @@ implementation: --task_config_uri, {inputValue: task_config_uri}, --start_at, {inputValue: start_at}, --resource_config_uri, {inputValue: resource_config_uri}, - --bootstrap_resource_config_hash, {inputValue: bootstrap_resource_config_hash}, --stop_after, {inputValue: stop_after}, --cpu_docker_uri, {inputValue: cpu_docker_uri}, --cuda_docker_uri, {inputValue: cuda_docker_uri}, diff --git a/gigl/orchestration/kubeflow/kfp_orchestrator.py b/gigl/orchestration/kubeflow/kfp_orchestrator.py index ebbdcc9f0..e9baa0a65 100644 --- a/gigl/orchestration/kubeflow/kfp_orchestrator.py +++ b/gigl/orchestration/kubeflow/kfp_orchestrator.py @@ -4,7 +4,6 @@ from pathlib import Path from typing import Optional, Union -import yaml from google.cloud import aiplatform from kfp.compiler import Compiler @@ -13,19 +12,15 @@ from gigl.common.logger import Logger from gigl.common.services.vertex_ai import VertexAIService from gigl.common.types.resource_config import CommonPipelineComponentConfigs -from gigl.common.utils.proto_utils import ProtoUtils, get_proto_fingerprint +from gigl.env.pipelines_config import get_resource_config 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() @@ -41,7 +36,6 @@ ) DEFAULT_START_AT_COMPONENT = "config_populator" -_CONFIG_RESOLUTION_PIPELINE_PARAMETER = "bootstrap_resource_config_hash" class KfpOrchestrator: @@ -147,43 +141,14 @@ def run( assert file_loader.does_uri_exist(compiled_pipeline_path), ( f"Compiled pipeline path {compiled_pipeline_path} does not exist." ) - compiled_pipeline_file = file_loader.load_to_temp_file( - file_uri_src=compiled_pipeline_path, - delete=True, - should_create_symlinks_if_possible=False, - ) - with open(compiled_pipeline_file.name, "r") as file: - compiled_pipeline = yaml.safe_load(file) - compiled_pipeline_file.close() - pipeline_parameters = ( - compiled_pipeline.get("root", {}) - .get("inputDefinitions", {}) - .get("parameters", {}) - if isinstance(compiled_pipeline, dict) - else {} - ) - if _CONFIG_RESOLUTION_PIPELINE_PARAMETER not in pipeline_parameters: - raise ValueError( - "The compiled pipeline predates config resolution support. " - "Recompile the pipeline before running it." - ) logger.info(f"Skipping pipeline compilation; will use {compiled_pipeline_path}") - resource_config = GiglResourceConfigWrapper( - ProtoUtils().read_proto_from_yaml( - uri=resource_config_uri, - proto_cls=GiglResourceConfig, - ) - ) - resolved_resource_config = resource_config.get_resolved_resource_config() + resource_config = get_resource_config(resource_config_uri=resource_config_uri) run_keyword_args = { "job_name": applied_task_identifier, "start_at": start_at, "template_or_frozen_config_uri": task_config_uri.uri, "resource_config_uri": resource_config_uri.uri, - "bootstrap_resource_config_hash": get_proto_fingerprint( - resolved_resource_config - ), } # We need to provide *some* notification emails, other wise the cleanup component will fail. # Ideally, we'd be able to provide None and have it handle it, but for whatever reason diff --git a/gigl/orchestration/kubeflow/kfp_pipeline.py b/gigl/orchestration/kubeflow/kfp_pipeline.py index 3f42a8c48..8ba3659c1 100644 --- a/gigl/orchestration/kubeflow/kfp_pipeline.py +++ b/gigl/orchestration/kubeflow/kfp_pipeline.py @@ -56,7 +56,6 @@ def _generate_component_task( common_pipeline_component_configs: CommonPipelineComponentConfigs, start_at: Optional[str] = None, stop_after: Optional[str] = None, - bootstrap_resource_config_hash: Optional[str] = None, ) -> PipelineTask: component_task: PipelineTask if component == GiGLComponents.ConfigPopulator: @@ -70,16 +69,11 @@ def _generate_component_task( ) elif component == GiGLComponents.ConfigValidator: - if bootstrap_resource_config_hash is None: - raise ValueError( - "ConfigValidator requires a bootstrap resource config hash." - ) component_task = _speced_component_op_dict[component]( job_name=job_name, task_config_uri=task_config_uri, start_at=start_at, resource_config_uri=resource_config_uri, - bootstrap_resource_config_hash=bootstrap_resource_config_hash, stop_after=stop_after, cpu_docker_uri=common_pipeline_component_configs.cpu_container_image, cuda_docker_uri=common_pipeline_component_configs.cuda_container_image, @@ -150,7 +144,6 @@ def _generate_component_tasks( job_name: str, template_or_frozen_config_uri: str, resource_config_uri: str, - bootstrap_resource_config_hash: str, common_pipeline_component_configs: CommonPipelineComponentConfigs, start_at: Optional[str] = None, stop_after: Optional[str] = None, @@ -162,7 +155,6 @@ def _generate_component_tasks( start_at=start_at, stop_after=stop_after, resource_config_uri=resource_config_uri, - bootstrap_resource_config_hash=bootstrap_resource_config_hash, common_pipeline_component_configs=common_pipeline_component_configs, ) validation_check_task.set_caching_options(enable_caching=False) @@ -267,7 +259,6 @@ def pipeline( job_name: str, template_or_frozen_config_uri: str, resource_config_uri: str, - bootstrap_resource_config_hash: str, start_at: str = GiGLComponents.ConfigPopulator.value, stop_after: Optional[str] = None, notification_emails: Optional[List[str]] = None, @@ -280,7 +271,6 @@ def pipeline( job_name=job_name, template_or_frozen_config_uri=template_or_frozen_config_uri, resource_config_uri=resource_config_uri, - bootstrap_resource_config_hash=bootstrap_resource_config_hash, common_pipeline_component_configs=common_pipeline_component_configs, start_at=start_at, stop_after=stop_after, diff --git a/gigl/src/validation_check/config_validator.py b/gigl/src/validation_check/config_validator.py index 092c065d6..53a185dc9 100644 --- a/gigl/src/validation_check/config_validator.py +++ b/gigl/src/validation_check/config_validator.py @@ -5,11 +5,7 @@ from gigl.common import GcsUri, Uri, UriFactory from gigl.common.logger import Logger from gigl.common.utils.gcs import GcsUtils -from gigl.common.utils.proto_utils import ( - ProtoUtils, - get_proto_fingerprint, - proto_to_yaml, -) +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 ( @@ -363,7 +359,6 @@ def _validate_resolved_configs( def resolve_configs( task_config_uri: Uri, resource_config_uri: Uri, - bootstrap_resource_config_hash: Optional[str] = None, ) -> tuple[gbml_config_pb2.GbmlConfig, GiglResourceConfig]: """Resolve task and resource configs into self-contained protobufs.""" proto_utils = ProtoUtils() @@ -378,15 +373,6 @@ def resolve_configs( ) ) resource_config = source_resource_config_wrapper.get_resolved_resource_config() - - resolved_resource_config_hash = get_proto_fingerprint(resource_config) - if ( - bootstrap_resource_config_hash is not None - and bootstrap_resource_config_hash != resolved_resource_config_hash - ): - raise ValueError( - "Resource config changed between pipeline submission and validation." - ) return task_config, resource_config @@ -396,12 +382,10 @@ def kfp_validation_checks( start_at: str, resource_config_uri: Uri, stop_after: Optional[str] = None, - bootstrap_resource_config_hash: 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, - bootstrap_resource_config_hash=bootstrap_resource_config_hash, ) should_use_live_sgs_backend = _validate_resolved_configs( job_name=job_name, @@ -436,11 +420,11 @@ def materialize_resolved_configs( ) task_config_uri = GcsUri.join( snapshot_root, - f"resolved_task_config_{get_proto_fingerprint(task_config)}.yaml", + "resolved_task_config.yaml", ) resource_config_uri = GcsUri.join( snapshot_root, - f"resolved_resource_config_{get_proto_fingerprint(resource_config)}.yaml", + "resolved_resource_config.yaml", ) gcs_utils = GcsUtils(project=resource_config_wrapper.project) gcs_utils.upload_from_string( @@ -489,12 +473,6 @@ def _write_kfp_output(path: str, value: str) -> None: type=str, help="Runtime argument for resource and env specifications of each component", ) - parser.add_argument( - "--bootstrap_resource_config_hash", - type=str, - required=True, - help="Resource config fingerprint computed during pipeline submission", - ) parser.add_argument( "--output_file_path_resolved_task_config_uri", type=str, @@ -534,7 +512,6 @@ def _write_kfp_output(path: str, value: str) -> None: task_config, resource_config = resolve_configs( task_config_uri=task_config_uri, resource_config_uri=resource_config_uri, - bootstrap_resource_config_hash=args.bootstrap_resource_config_hash, ) ( resolved_task_config_uri, diff --git a/tests/unit/common/utils/proto_utils_test.py b/tests/unit/common/utils/proto_utils_test.py index e7d0d852b..d9e4bd473 100644 --- a/tests/unit/common/utils/proto_utils_test.py +++ b/tests/unit/common/utils/proto_utils_test.py @@ -7,12 +7,9 @@ from absl.testing import absltest from hydra import initialize_config_dir -from hydra.errors import ConfigCompositionException -from omegaconf import OmegaConf from gigl.common import GcsUri, LocalUri from gigl.common.logger import Logger -from gigl.common.omegaconf_resolvers import now_resolver 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 @@ -145,7 +142,7 @@ def test_can_compose_resource_config_from_primary_parent(self): "example-project", ) - def test_resource_config_rejects_interpolation(self): + def test_resource_config_resolves_dynamic_interpolation(self): with TemporaryDirectory() as temp_directory: config_path = Path(temp_directory) / "resource.yaml" config_path.write_text( @@ -154,12 +151,17 @@ def test_resource_config_rejects_interpolation(self): ' project: "${oc.env:PROJECT_ID}"\n' ) - with self.assertRaises(ValueError): - self.proto_utils.read_proto_from_yaml( + 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" @@ -207,7 +209,7 @@ def test_resource_composition_ignores_dynamic_resolver_in_unselected_file(self): "example-project", ) - def test_resource_composition_rejects_dynamic_resolver_in_selected_file(self): + 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() @@ -219,13 +221,18 @@ def test_resource_composition_rejects_dynamic_resolver_in_selected_file(self): 'common_compute_config:\n project: "${oc.env:PROJECT_ID}"\n' ) - with self.assertRaises(ValueError): - self.proto_utils.read_proto_from_yaml( + 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, ) - def test_resource_composition_rejects_dynamic_nested_default(self): + 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() @@ -236,7 +243,7 @@ def test_resource_composition_rejects_dynamic_nested_default(self): ) (config_root / "base" / "resource.yaml").write_text( "defaults:\n" - " - compute@shared_resource_config: " + " - /compute@shared_resource_config: " "${oc.env:RESOURCE_PROFILE}\n" " - _self_\n" ) @@ -246,15 +253,17 @@ def test_resource_composition_rejects_dynamic_nested_default(self): " region: us-central1\n" ) - with ( - patch.dict(os.environ, {"RESOURCE_PROFILE": "local"}), - self.assertRaises(ConfigCompositionException), - ): - self.proto_utils.read_proto_from_yaml( + with patch.dict(os.environ, {"RESOURCE_PROFILE": "local"}): + 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_composed_primary_requires_yaml_extension(self): with TemporaryDirectory() as temp_directory: config_path = Path(temp_directory) / "task.yml" @@ -316,24 +325,6 @@ def test_composition_preserves_foreign_hydra_context(self): proto_cls=gbml_config_pb2.GbmlConfig, ) - def test_composition_restores_foreign_omegaconf_resolver(self): - with TemporaryDirectory() as temp_directory: - config_path = Path(temp_directory) / "task.yaml" - config_path.write_text("defaults:\n - _self_\n") - OmegaConf.register_new_resolver( - "now", lambda *_: "foreign-now", replace=True - ) - try: - self.proto_utils.read_proto_from_yaml( - uri=LocalUri(config_path), - proto_cls=gbml_config_pb2.GbmlConfig, - ) - - config = OmegaConf.create({"value": "${now:any-format}"}) - self.assertEqual(config.value, "foreign-now") - finally: - OmegaConf.register_new_resolver("now", now_resolver, replace=True) - 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 b5dec6670..922f9c550 100644 --- a/tests/unit/orchestration/kubeflow/component_spec_test.py +++ b/tests/unit/orchestration/kubeflow/component_spec_test.py @@ -16,7 +16,6 @@ class ComponentSpecTest(TestCase): def test_config_validator_resolves_both_configs(self) -> None: component_spec = _read_component_spec("config_validator") - self.assertIn("name: bootstrap_resource_config_hash", component_spec) 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) diff --git a/tests/unit/orchestration/kubeflow/kfp_pipeline_test.py b/tests/unit/orchestration/kubeflow/kfp_pipeline_test.py index 9d2757793..5d4652816 100644 --- a/tests/unit/orchestration/kubeflow/kfp_pipeline_test.py +++ b/tests/unit/orchestration/kubeflow/kfp_pipeline_test.py @@ -25,10 +25,6 @@ def test_validator_outputs_are_pipeline_inputs_for_downstream_tasks(self) -> Non pipeline_text = output_path.read_text() pipeline_spec = yaml.safe_load(pipeline_text) - self.assertIn( - "bootstrap_resource_config_hash", - pipeline_spec["root"]["inputDefinitions"]["parameters"], - ) validator_task = pipeline_spec["components"]["comp-exit-handler-1"]["dag"][ "tasks" ]["kfp-validation-check"] diff --git a/tests/unit/src/validation/config_validator_test.py b/tests/unit/src/validation/config_validator_test.py index f95750384..cae63b691 100644 --- a/tests/unit/src/validation/config_validator_test.py +++ b/tests/unit/src/validation/config_validator_test.py @@ -9,7 +9,7 @@ import gigl.env.pipelines_config from gigl.common import Uri, UriFactory -from gigl.common.utils.proto_utils import ProtoUtils, get_proto_fingerprint +from gigl.common.utils.proto_utils import ProtoUtils from gigl.src.validation_check.config_validator import ( kfp_validation_checks, materialize_resolved_configs, @@ -355,36 +355,6 @@ def test_resource_config_validation_failure_with_mock_configs( resource_config_uri=self._live_resource_config_uri, ) - def test_resource_config_bootstrap_hash_must_match(self) -> None: - with self.assertRaises(ValueError): - kfp_validation_checks( - job_name="resource_config_validation_test", - task_config_uri=self._live_task_config_uri, - start_at="config_populator", - resource_config_uri=self._live_resource_config_uri, - bootstrap_resource_config_hash="stale-resource-config", - ) - - def test_resource_config_bootstrap_hash_accepts_resolved_config(self) -> None: - resource_config = self._proto_utils.read_proto_from_yaml( - uri=self._live_resource_config_uri, - proto_cls=gigl_resource_config_pb2.GiglResourceConfig, - ) - - _, resolved_resource_config, should_use_glt_backend = kfp_validation_checks( - job_name="resource_config_validation_test", - task_config_uri=self._live_task_config_uri, - start_at="config_populator", - resource_config_uri=self._live_resource_config_uri, - bootstrap_resource_config_hash=get_proto_fingerprint(resource_config), - ) - - self.assertTrue(should_use_glt_backend) - self.assertEqual( - get_proto_fingerprint(resolved_resource_config), - get_proto_fingerprint(resource_config), - ) - @patch("gigl.src.validation_check.config_validator.GcsUtils") def test_materializes_both_resolved_configs(self, mock_gcs_utils) -> None: task_config = _create_valid_live_subgraph_sampling_task_config() @@ -396,15 +366,15 @@ def test_materializes_both_resolved_configs(self, mock_gcs_utils) -> None: resource_config=resource_config, ) - self.assertTrue( - task_uri.uri.startswith( - "gs://test-temp-regional/config-resolution-test/config_validator/" - ) + self.assertEqual( + task_uri.uri, + "gs://test-temp-regional/config-resolution-test/config_validator/" + "resolved_task_config.yaml", ) - self.assertTrue( - resource_uri.uri.startswith( - "gs://test-temp-regional/config-resolution-test/config_validator/" - ) + self.assertEqual( + resource_uri.uri, + "gs://test-temp-regional/config-resolution-test/config_validator/" + "resolved_resource_config.yaml", ) self.assertEqual( mock_gcs_utils.return_value.upload_from_string.call_count, From b76325744919acdbb7577e48a884efea340beb40 Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Thu, 30 Jul 2026 17:03:21 +0000 Subject: [PATCH 4/9] Simplify Hydra config opt-in and resolution --- .../config_guides/hydra_composition.md | 4 +- gigl/common/utils/proto_utils.py | 12 +-- .../types/pb_wrappers/gigl_resource_config.py | 16 ---- gigl/src/validation_check/config_validator.py | 12 +-- tests/config_tests/notebooks_test.py | 5 +- tests/unit/common/utils/proto_utils_test.py | 40 ++++++++- .../src/validation/config_validator_test.py | 85 +++++++++++++++---- 7 files changed, 122 insertions(+), 52 deletions(-) diff --git a/docs/user_guide/config_guides/hydra_composition.md b/docs/user_guide/config_guides/hydra_composition.md index d837ae392..2163dcd19 100644 --- a/docs/user_guide/config_guides/hydra_composition.md +++ b/docs/user_guide/config_guides/hydra_composition.md @@ -1,7 +1,7 @@ # Composing task and resource configs with Hydra -GiGL task and resource configs can opt into [Hydra Defaults Lists](https://hydra.cc/docs/advanced/defaults_list/) by -adding a top-level `defaults` list. Configs without `defaults` continue through the existing YAML and OmegaConf loader. +GiGL protobuf YAML configs can opt into [Hydra Defaults Lists](https://hydra.cc/docs/advanced/defaults_list/) by adding +a top-level `defaults` list. Configs without `defaults` continue through the existing YAML and OmegaConf loader. ## Config root diff --git a/gigl/common/utils/proto_utils.py b/gigl/common/utils/proto_utils.py index 9037b320d..9b8f09161 100644 --- a/gigl/common/utils/proto_utils.py +++ b/gigl/common/utils/proto_utils.py @@ -16,11 +16,6 @@ T = TypeVar("T", bound=message.Message) -_HYDRA_PROTO_TYPES = { - "snapchat.research.gbml.GbmlConfig", - "snapchat.research.gbml.GiglResourceConfig", -} - def proto_to_yaml(proto: message.Message) -> str: """Serialize a protobuf message to canonical YAML. @@ -45,13 +40,8 @@ def read_proto_from_yaml(self, uri: Uri, proto_cls: Type[T]) -> T: with open(tfh.name, "r") as file: raw_data = yaml.safe_load(file) tfh.close() - proto_type = proto_cls.DESCRIPTOR.full_name - if ( - isinstance(raw_data, dict) - and "defaults" in raw_data - and proto_type in _HYDRA_PROTO_TYPES - ): + if isinstance(raw_data, dict) and "defaults" in raw_data: obj_dict = compose_yaml_config(uri=uri) else: omega_conf_obj = OmegaConf.create(raw_data) diff --git a/gigl/src/common/types/pb_wrappers/gigl_resource_config.py b/gigl/src/common/types/pb_wrappers/gigl_resource_config.py index a2809c927..6a5d0b40b 100644 --- a/gigl/src/common/types/pb_wrappers/gigl_resource_config.py +++ b/gigl/src/common/types/pb_wrappers/gigl_resource_config.py @@ -97,22 +97,6 @@ def shared_resource_config(self) -> SharedResourceConfig: return self._loaded_shared_resource_config - def get_resolved_resource_config(self) -> GiglResourceConfig: - """Return a self-contained resource config. - - If the source config references a shared resource config URI, the - referenced protobuf is copied into the returned message. - - Returns: - A new resource config with an inline shared resource config. - """ - resolved_resource_config = GiglResourceConfig() - resolved_resource_config.CopyFrom(self.resource_config) - resolved_resource_config.shared_resource_config.CopyFrom( - self.shared_resource_config - ) - return resolved_resource_config - @property def get_resource_config_uri(self) -> str: parser = argparse.ArgumentParser() diff --git a/gigl/src/validation_check/config_validator.py b/gigl/src/validation_check/config_validator.py index 53a185dc9..aabd8daf6 100644 --- a/gigl/src/validation_check/config_validator.py +++ b/gigl/src/validation_check/config_validator.py @@ -366,13 +366,13 @@ def resolve_configs( uri=task_config_uri, proto_cls=gbml_config_pb2.GbmlConfig, ) - source_resource_config_wrapper = GiglResourceConfigWrapper( - proto_utils.read_proto_from_yaml( - uri=resource_config_uri, - proto_cls=GiglResourceConfig, - ) + resource_config = proto_utils.read_proto_from_yaml( + uri=resource_config_uri, + proto_cls=GiglResourceConfig, + ) + resource_config.shared_resource_config.CopyFrom( + GiglResourceConfigWrapper(resource_config).shared_resource_config ) - resource_config = source_resource_config_wrapper.get_resolved_resource_config() return task_config, resource_config diff --git a/tests/config_tests/notebooks_test.py b/tests/config_tests/notebooks_test.py index 7dd1df456..098f9c548 100644 --- a/tests/config_tests/notebooks_test.py +++ b/tests/config_tests/notebooks_test.py @@ -65,8 +65,11 @@ def setUp(self): logger.info(f"Using GCS 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.get_resolved_resource_config(), + proto=resource_config.resource_config, uri=gcs_uri, ) self._notebooks = [ diff --git a/tests/unit/common/utils/proto_utils_test.py b/tests/unit/common/utils/proto_utils_test.py index d9e4bd473..643bd7145 100644 --- a/tests/unit/common/utils/proto_utils_test.py +++ b/tests/unit/common/utils/proto_utils_test.py @@ -7,13 +7,18 @@ from absl.testing import absltest from hydra import initialize_config_dir +from hydra.errors import MissingConfigException +from omegaconf import OmegaConf from gigl.common import GcsUri, 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 +from snapchat.research.gbml.gigl_resource_config_pb2 import ( + GiglResourceConfig, + SharedResourceConfig, +) from tests.test_assets.test_case import TestCase logger = Logger() @@ -142,6 +147,28 @@ def test_can_compose_resource_config_from_primary_parent(self): "example-project", ) + def test_defaults_enable_composition_for_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.read_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" @@ -325,6 +352,17 @@ def test_composition_preserves_foreign_hydra_context(self): 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/src/validation/config_validator_test.py b/tests/unit/src/validation/config_validator_test.py index cae63b691..995911b3e 100644 --- a/tests/unit/src/validation/config_validator_test.py +++ b/tests/unit/src/validation/config_validator_test.py @@ -1,6 +1,7 @@ import os import shutil import tempfile +from pathlib import Path from unittest.mock import patch import google.protobuf.message @@ -13,6 +14,7 @@ from gigl.src.validation_check.config_validator import ( kfp_validation_checks, materialize_resolved_configs, + resolve_configs, ) from snapchat.research.gbml import ( gbml_config_pb2, @@ -356,30 +358,83 @@ def test_resource_config_validation_failure_with_mock_configs( ) @patch("gigl.src.validation_check.config_validator.GcsUtils") - def test_materializes_both_resolved_configs(self, mock_gcs_utils) -> None: - task_config = _create_valid_live_subgraph_sampling_task_config() - resource_config = _create_valid_live_subgraph_sampling_resource_config() + 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.uri, - "gs://test-temp-regional/config-resolution-test/config_validator/" - "resolved_task_config.yaml", - ) - self.assertEqual( - resource_uri.uri, - "gs://test-temp-regional/config-resolution-test/config_validator/" - "resolved_resource_config.yaml", + 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, ) - self.assertEqual( - mock_gcs_utils.return_value.upload_from_string.call_count, - 2, + 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) if __name__ == "__main__": From ea4742e69b62dce33f25aa9c85b7a852a32af7ae Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Thu, 30 Jul 2026 17:17:09 +0000 Subject: [PATCH 5/9] Make Hydra composition explicit --- .../config_guides/hydra_composition.md | 14 ++++---- gigl/common/utils/proto_utils.py | 14 ++++---- .../kubeflow/kfp_orchestrator.py | 19 ++++++++++- gigl/src/validation_check/config_validator.py | 34 ++++++++++++++----- .../config_tests/assert_yaml_configs_parse.py | 8 ++--- tests/config_tests/notebooks_test.py | 22 ++++++++++-- tests/unit/common/utils/proto_utils_test.py | 29 ++++++++++------ 7 files changed, 100 insertions(+), 40 deletions(-) diff --git a/docs/user_guide/config_guides/hydra_composition.md b/docs/user_guide/config_guides/hydra_composition.md index 2163dcd19..bd1ec46a7 100644 --- a/docs/user_guide/config_guides/hydra_composition.md +++ b/docs/user_guide/config_guides/hydra_composition.md @@ -1,7 +1,9 @@ # Composing task and resource configs with Hydra -GiGL protobuf YAML configs can opt into [Hydra Defaults Lists](https://hydra.cc/docs/advanced/defaults_list/) by adding -a top-level `defaults` list. Configs without `defaults` continue through the existing YAML and OmegaConf loader. +Local 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 these source configs before publishing plain YAML snapshots. +`ProtoUtils.read_proto_from_yaml` remains the URI-agnostic, non-composing reader; `ProtoUtils.compose_proto_from_yaml` +is the explicit local-only composition API. ## Config root @@ -59,10 +61,10 @@ common_compute_config: temp_regional_assets_bucket: gs://example-bucket ``` -Resource config is read once before pipeline submission to select the project, region, service account, and staging -bucket. ConfigValidator resolves 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 read. +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 diff --git a/gigl/common/utils/proto_utils.py b/gigl/common/utils/proto_utils.py index 9b8f09161..dc144cc14 100644 --- a/gigl/common/utils/proto_utils.py +++ b/gigl/common/utils/proto_utils.py @@ -39,13 +39,9 @@ def read_proto_from_yaml(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, "r") as file: raw_data = yaml.safe_load(file) - tfh.close() - - if isinstance(raw_data, dict) and "defaults" in raw_data: - obj_dict = compose_yaml_config(uri=uri) - else: omega_conf_obj = OmegaConf.create(raw_data) - obj_dict = OmegaConf.to_object(omega_conf_obj) + tfh.close() + obj_dict = OmegaConf.to_object(omega_conf_obj) if not isinstance(obj_dict, dict): raise TypeError( f"ProtoUtils.read_proto_from_yaml expected a mapping at the YAML root for " @@ -54,6 +50,12 @@ 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: LocalUri, proto_cls: Type[T]) -> T: + """Compose a local YAML config with Hydra and parse it as a protobuf.""" + obj_dict = compose_yaml_config(uri=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/kfp_orchestrator.py b/gigl/orchestration/kubeflow/kfp_orchestrator.py index e9baa0a65..26365858a 100644 --- a/gigl/orchestration/kubeflow/kfp_orchestrator.py +++ b/gigl/orchestration/kubeflow/kfp_orchestrator.py @@ -12,15 +12,20 @@ from gigl.common.logger import Logger from gigl.common.services.vertex_ai import VertexAIService from gigl.common.types.resource_config import CommonPipelineComponentConfigs +from gigl.common.utils.proto_utils import ProtoUtils from gigl.env.pipelines_config import get_resource_config 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,7 +148,19 @@ def run( ) logger.info(f"Skipping pipeline compilation; will use {compiled_pipeline_path}") - resource_config = get_resource_config(resource_config_uri=resource_config_uri) + if isinstance( + resource_config_uri, LocalUri + ) and resource_config_uri.get_basename().endswith(".yaml"): + resource_config = GiglResourceConfigWrapper( + ProtoUtils().compose_proto_from_yaml( + uri=resource_config_uri, + proto_cls=GiglResourceConfig, + ) + ) + else: + resource_config = get_resource_config( + resource_config_uri=resource_config_uri + ) run_keyword_args = { "job_name": applied_task_identifier, "start_at": start_at, diff --git a/gigl/src/validation_check/config_validator.py b/gigl/src/validation_check/config_validator.py index aabd8daf6..3ec824a04 100644 --- a/gigl/src/validation_check/config_validator.py +++ b/gigl/src/validation_check/config_validator.py @@ -2,7 +2,7 @@ from pathlib import Path from typing import Optional -from gigl.common import GcsUri, Uri, UriFactory +from gigl.common import GcsUri, LocalUri, Uri, UriFactory from gigl.common.logger import Logger from gigl.common.utils.gcs import GcsUtils from gigl.common.utils.proto_utils import ProtoUtils, proto_to_yaml @@ -362,14 +362,30 @@ def resolve_configs( ) -> tuple[gbml_config_pb2.GbmlConfig, GiglResourceConfig]: """Resolve task and resource configs into self-contained protobufs.""" proto_utils = ProtoUtils() - task_config = proto_utils.read_proto_from_yaml( - uri=task_config_uri, - proto_cls=gbml_config_pb2.GbmlConfig, - ) - resource_config = proto_utils.read_proto_from_yaml( - uri=resource_config_uri, - proto_cls=GiglResourceConfig, - ) + if isinstance( + task_config_uri, LocalUri + ) and task_config_uri.get_basename().endswith(".yaml"): + task_config = proto_utils.compose_proto_from_yaml( + uri=task_config_uri, + proto_cls=gbml_config_pb2.GbmlConfig, + ) + else: + task_config = proto_utils.read_proto_from_yaml( + uri=task_config_uri, + proto_cls=gbml_config_pb2.GbmlConfig, + ) + if isinstance( + resource_config_uri, LocalUri + ) and resource_config_uri.get_basename().endswith(".yaml"): + resource_config = proto_utils.compose_proto_from_yaml( + uri=resource_config_uri, + proto_cls=GiglResourceConfig, + ) + else: + resource_config = proto_utils.read_proto_from_yaml( + uri=resource_config_uri, + proto_cls=GiglResourceConfig, + ) resource_config.shared_resource_config.CopyFrom( GiglResourceConfigWrapper(resource_config).shared_resource_config ) 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 098f9c548..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.types.pb_wrappers.gigl_resource_config import ( + GiglResourceConfigWrapper, +) +from snapchat.research.gbml.gigl_resource_config_pb2 import GiglResourceConfig logger = Logger() @@ -52,9 +57,20 @@ def setUp(self): logger.info(f"Using resource config URI: {resource_config_uri}") # Materialize the resource config in GCS so every machine running # notebook business logic can access the same self-contained snapshot. - resource_config = get_resource_config( - resource_config_uri=UriFactory.create_uri(resource_config_uri) - ) + 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( resource_config.temp_assets_regional_bucket_path, "testing", diff --git a/tests/unit/common/utils/proto_utils_test.py b/tests/unit/common/utils/proto_utils_test.py index 643bd7145..c21186442 100644 --- a/tests/unit/common/utils/proto_utils_test.py +++ b/tests/unit/common/utils/proto_utils_test.py @@ -3,6 +3,7 @@ from datetime import datetime, timedelta from pathlib import Path from tempfile import NamedTemporaryFile, TemporaryDirectory +from typing import cast from unittest.mock import patch from absl.testing import absltest @@ -110,7 +111,7 @@ def test_can_compose_task_config_from_primary_parent(self): "isGraphDirected: true\n" ) - task_config = self.proto_utils.read_proto_from_yaml( + task_config = self.proto_utils.compose_proto_from_yaml( uri=LocalUri(config_root / "task.yaml"), proto_cls=gbml_config_pb2.GbmlConfig, ) @@ -137,7 +138,7 @@ def test_can_compose_resource_config_from_primary_parent(self): " temp_regional_assets_bucket: gs://example-bucket\n" ) - resource_config = self.proto_utils.read_proto_from_yaml( + resource_config = self.proto_utils.compose_proto_from_yaml( uri=LocalUri(config_root / "resource.yaml"), proto_cls=GiglResourceConfig, ) @@ -147,7 +148,7 @@ def test_can_compose_resource_config_from_primary_parent(self): "example-project", ) - def test_defaults_enable_composition_for_any_proto(self): + def test_can_compose_any_proto(self): with TemporaryDirectory() as temp_directory: config_root = Path(temp_directory) (config_root / "compute").mkdir() @@ -159,7 +160,7 @@ def test_defaults_enable_composition_for_any_proto(self): "project: example-project\nregion: us-central1\n" ) - shared_resource_config = self.proto_utils.read_proto_from_yaml( + shared_resource_config = self.proto_utils.compose_proto_from_yaml( uri=LocalUri(config_path), proto_cls=SharedResourceConfig, ) @@ -226,7 +227,7 @@ def test_resource_composition_ignores_dynamic_resolver_in_unselected_file(self): 'run_name: "${now:%Y%m%d}"\n' ) - resource_config = self.proto_utils.read_proto_from_yaml( + resource_config = self.proto_utils.compose_proto_from_yaml( uri=LocalUri(config_path), proto_cls=GiglResourceConfig, ) @@ -249,7 +250,7 @@ def test_resource_composition_resolves_dynamic_resolver_in_selected_file(self): ) with patch.dict(os.environ, {"PROJECT_ID": "example-project"}): - resource_config = self.proto_utils.read_proto_from_yaml( + resource_config = self.proto_utils.compose_proto_from_yaml( uri=LocalUri(config_path), proto_cls=GiglResourceConfig, ) @@ -281,7 +282,7 @@ def test_resource_composition_resolves_dynamic_nested_default(self): ) with patch.dict(os.environ, {"RESOURCE_PROFILE": "local"}): - resource_config = self.proto_utils.read_proto_from_yaml( + resource_config = self.proto_utils.compose_proto_from_yaml( uri=LocalUri(config_path), proto_cls=GiglResourceConfig, ) @@ -297,7 +298,7 @@ def test_composed_primary_requires_yaml_extension(self): config_path.write_text("defaults:\n - _self_\n") with self.assertRaises(ValueError): - self.proto_utils.read_proto_from_yaml( + self.proto_utils.compose_proto_from_yaml( uri=LocalUri(config_path), proto_cls=gbml_config_pb2.GbmlConfig, ) @@ -307,7 +308,13 @@ def test_gcs_composition_is_not_supported(self): TypeError, "Hydra composition is not supported for GcsUri", ): - compose_yaml_config(GcsUri("gs://example-bucket/configs/task.yaml")) + self.proto_utils.compose_proto_from_yaml( + cast( + LocalUri, + GcsUri("gs://example-bucket/configs/task.yaml"), + ), + gbml_config_pb2.GbmlConfig, + ) def test_composition_is_thread_safe(self): with TemporaryDirectory() as temp_directory: @@ -324,7 +331,7 @@ def test_composition_is_thread_safe(self): def read_config(_: int) -> bool: return ( ProtoUtils() - .read_proto_from_yaml( + .compose_proto_from_yaml( uri=LocalUri(config_path), proto_cls=gbml_config_pb2.GbmlConfig, ) @@ -347,7 +354,7 @@ def test_composition_preserves_foreign_hydra_context(self): version_base="1.3", ): with self.assertRaises(RuntimeError): - self.proto_utils.read_proto_from_yaml( + self.proto_utils.compose_proto_from_yaml( uri=LocalUri(config_path), proto_cls=gbml_config_pb2.GbmlConfig, ) From 7451fdac0974c83e7f069a98a46a6f63a53a551e Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Thu, 30 Jul 2026 18:48:52 +0000 Subject: [PATCH 6/9] Always compose configs in ConfigValidator --- .../config_guides/hydra_composition.md | 12 ++-- gigl/common/omegaconf_resolvers.py | 13 ++-- gigl/common/utils/hydra_config.py | 25 +++++--- .../kubeflow/kfp_orchestrator.py | 4 +- gigl/src/validation_check/config_validator.py | 63 +++++++++++-------- tests/unit/common/utils/proto_utils_test.py | 16 +---- .../src/validation/config_validator_test.py | 36 ++++++++++- 7 files changed, 102 insertions(+), 67 deletions(-) diff --git a/docs/user_guide/config_guides/hydra_composition.md b/docs/user_guide/config_guides/hydra_composition.md index bd1ec46a7..e7175e905 100644 --- a/docs/user_guide/config_guides/hydra_composition.md +++ b/docs/user_guide/config_guides/hydra_composition.md @@ -1,14 +1,15 @@ # Composing task and resource configs with Hydra -Local 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 these source configs before publishing plain YAML snapshots. +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 explicit local-only composition API. ## Config root The primary file's parent directory is its Hydra config root. Composed primary files and selected fragments must use -Hydra's `.yaml` extension. Plain legacy `.yml` files without `defaults` remain supported. +Hydra's `.yaml` extension. Plain legacy `.yml` files remain supported by the non-composing reader, but ConfigValidator +inputs must use `.yaml`. ```text configs/ @@ -21,8 +22,9 @@ configs/ ``` Keep primary files directly in the bundle root. GiGL does not search for a repository root or a directory named -`configs`. Composition currently requires a local primary file. GCS and HTTP primaries remain supported only as plain, -single-file YAML without a Defaults List. +`configs`. Defaults List composition currently requires a local config bundle. ConfigValidator downloads GCS and HTTP +primaries to a temporary local file before composing them, so remote primaries remain supported only as plain, +single-file YAML. ## Task config example diff --git a/gigl/common/omegaconf_resolvers.py b/gigl/common/omegaconf_resolvers.py index 841d061c1..6653620e5 100644 --- a/gigl/common/omegaconf_resolvers.py +++ b/gigl/common/omegaconf_resolvers.py @@ -157,26 +157,23 @@ def git_hash_resolver() -> str: return "" -def register_resolvers(replace: bool = False) -> None: +def register_resolvers() -> None: """Register all custom OmegaConf resolvers. This function should be called at application startup to register all custom resolvers with OmegaConf. - - Args: - replace: Whether to replace resolvers registered by another library. """ - if replace or not OmegaConf.has_resolver("now"): + if not OmegaConf.has_resolver("now"): logger.info("Registering OmegaConf resolver 'now'") - OmegaConf.register_new_resolver("now", now_resolver, replace=replace) + OmegaConf.register_new_resolver("now", now_resolver) else: logger.debug( "OmegaConf resolver 'now' already registered, skipping registration" ) - if replace or not OmegaConf.has_resolver("git_hash"): + if not OmegaConf.has_resolver("git_hash"): logger.info("Registering OmegaConf resolver 'git_hash'") - OmegaConf.register_new_resolver("git_hash", git_hash_resolver, replace=replace) + OmegaConf.register_new_resolver("git_hash", git_hash_resolver) else: logger.debug( "OmegaConf resolver 'git_hash' already registered, skipping registration" diff --git a/gigl/common/utils/hydra_config.py b/gigl/common/utils/hydra_config.py index 03b32c1d7..44213a353 100644 --- a/gigl/common/utils/hydra_config.py +++ b/gigl/common/utils/hydra_config.py @@ -11,13 +11,21 @@ from hydra.core.global_hydra import GlobalHydra from omegaconf import OmegaConf -from gigl.common import LocalUri, Uri -from gigl.common.omegaconf_resolvers import register_resolvers +from gigl.common import LocalUri +from gigl.common.omegaconf_resolvers import now_resolver, register_resolvers _COMPOSE_LOCK = threading.RLock() -def compose_yaml_config(uri: Uri) -> dict[str, Any]: +def _register_gigl_now_resolver() -> None: + OmegaConf.register_new_resolver( + "now", + now_resolver, + replace=True, + ) + + +def compose_yaml_config(uri: LocalUri) -> dict[str, Any]: """Compose a YAML config with Hydra using its parent as the config root. Args: @@ -27,7 +35,6 @@ def compose_yaml_config(uri: Uri) -> dict[str, Any]: A fully composed and resolved mapping. Raises: - TypeError: If composition is requested for an unsupported URI. ValueError: If the primary filename is unsupported or the result is not a mapping. RuntimeError: If another Hydra application owns the global context. @@ -37,10 +44,7 @@ def compose_yaml_config(uri: Uri) -> dict[str, Any]: raise ValueError(f"Hydra primary config must use the .yaml extension: {uri}") config_name = primary_name.rsplit(".", 1)[0] - if isinstance(uri, LocalUri): - config_root = Path(uri.uri).absolute().parent - else: - raise TypeError(f"Hydra composition is not supported for {type(uri).__name__}.") + config_root = Path(uri.uri).absolute().parent with _COMPOSE_LOCK: if GlobalHydra.instance().is_initialized(): @@ -56,11 +60,12 @@ def compose_yaml_config(uri: Uri) -> dict[str, Any]: # Hydra installs a one-argument ``now`` resolver during # initialization. GiGL's resolver is a backward-compatible # superset that also supports offsets. - register_resolvers(replace=True) + register_resolvers() + _register_gigl_now_resolver() composed = compose(config_name=config_name, overrides=[]) resolved = OmegaConf.to_container(composed, resolve=True) finally: - register_resolvers(replace=True) + _register_gigl_now_resolver() if not isinstance(resolved, dict): raise ValueError( diff --git a/gigl/orchestration/kubeflow/kfp_orchestrator.py b/gigl/orchestration/kubeflow/kfp_orchestrator.py index 26365858a..1f7230a99 100644 --- a/gigl/orchestration/kubeflow/kfp_orchestrator.py +++ b/gigl/orchestration/kubeflow/kfp_orchestrator.py @@ -148,9 +148,7 @@ def run( ) logger.info(f"Skipping pipeline compilation; will use {compiled_pipeline_path}") - if isinstance( - resource_config_uri, LocalUri - ) and resource_config_uri.get_basename().endswith(".yaml"): + if isinstance(resource_config_uri, LocalUri): resource_config = GiglResourceConfigWrapper( ProtoUtils().compose_proto_from_yaml( uri=resource_config_uri, diff --git a/gigl/src/validation_check/config_validator.py b/gigl/src/validation_check/config_validator.py index 3ec824a04..8d7759185 100644 --- a/gigl/src/validation_check/config_validator.py +++ b/gigl/src/validation_check/config_validator.py @@ -1,6 +1,9 @@ import argparse from pathlib import Path -from typing import Optional +from tempfile import NamedTemporaryFile +from typing import Optional, Type, TypeVar + +from google.protobuf import message from gigl.common import GcsUri, LocalUri, Uri, UriFactory from gigl.common.logger import Logger @@ -11,6 +14,7 @@ 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.gigl_runtime import initialize_gigl_runtime from gigl.src.validation_check.libs.frozen_config_path_checks import ( assert_preprocessed_metadata_exists, @@ -198,6 +202,8 @@ logger = Logger() +T = TypeVar("T", bound=message.Message) + # 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 @@ -362,36 +368,43 @@ def resolve_configs( ) -> tuple[gbml_config_pb2.GbmlConfig, GiglResourceConfig]: """Resolve task and resource configs into self-contained protobufs.""" proto_utils = ProtoUtils() - if isinstance( - task_config_uri, LocalUri - ) and task_config_uri.get_basename().endswith(".yaml"): - task_config = proto_utils.compose_proto_from_yaml( - uri=task_config_uri, - proto_cls=gbml_config_pb2.GbmlConfig, - ) - else: - task_config = proto_utils.read_proto_from_yaml( - uri=task_config_uri, - proto_cls=gbml_config_pb2.GbmlConfig, - ) - if isinstance( - resource_config_uri, LocalUri - ) and resource_config_uri.get_basename().endswith(".yaml"): - resource_config = proto_utils.compose_proto_from_yaml( - uri=resource_config_uri, - proto_cls=GiglResourceConfig, - ) - else: - resource_config = proto_utils.read_proto_from_yaml( - uri=resource_config_uri, - proto_cls=GiglResourceConfig, - ) + task_config = _compose_proto_from_uri( + proto_utils=proto_utils, + uri=task_config_uri, + proto_cls=gbml_config_pb2.GbmlConfig, + ) + resource_config = _compose_proto_from_uri( + proto_utils=proto_utils, + uri=resource_config_uri, + proto_cls=GiglResourceConfig, + ) resource_config.shared_resource_config.CopyFrom( GiglResourceConfigWrapper(resource_config).shared_resource_config ) return task_config, resource_config +def _compose_proto_from_uri( + proto_utils: ProtoUtils, + uri: Uri, + proto_cls: Type[T], +) -> T: + if isinstance(uri, LocalUri): + return proto_utils.compose_proto_from_yaml(uri=uri, proto_cls=proto_cls) + + with NamedTemporaryFile(suffix=".yaml") as temp_file: + local_uri = LocalUri(temp_file.name) + FileLoader().load_file( + file_uri_src=uri, + file_uri_dst=local_uri, + should_create_symlinks_if_possible=False, + ) + return proto_utils.compose_proto_from_yaml( + uri=local_uri, + proto_cls=proto_cls, + ) + + def kfp_validation_checks( job_name: str, task_config_uri: Uri, diff --git a/tests/unit/common/utils/proto_utils_test.py b/tests/unit/common/utils/proto_utils_test.py index c21186442..eaa4ae67a 100644 --- a/tests/unit/common/utils/proto_utils_test.py +++ b/tests/unit/common/utils/proto_utils_test.py @@ -3,7 +3,6 @@ from datetime import datetime, timedelta from pathlib import Path from tempfile import NamedTemporaryFile, TemporaryDirectory -from typing import cast from unittest.mock import patch from absl.testing import absltest @@ -11,7 +10,7 @@ from hydra.errors import MissingConfigException from omegaconf import OmegaConf -from gigl.common import GcsUri, LocalUri +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 @@ -303,19 +302,6 @@ def test_composed_primary_requires_yaml_extension(self): proto_cls=gbml_config_pb2.GbmlConfig, ) - def test_gcs_composition_is_not_supported(self): - with self.assertRaisesRegex( - TypeError, - "Hydra composition is not supported for GcsUri", - ): - self.proto_utils.compose_proto_from_yaml( - cast( - LocalUri, - GcsUri("gs://example-bucket/configs/task.yaml"), - ), - gbml_config_pb2.GbmlConfig, - ) - def test_composition_is_thread_safe(self): with TemporaryDirectory() as temp_directory: config_root = Path(temp_directory) diff --git a/tests/unit/src/validation/config_validator_test.py b/tests/unit/src/validation/config_validator_test.py index 995911b3e..992760b8a 100644 --- a/tests/unit/src/validation/config_validator_test.py +++ b/tests/unit/src/validation/config_validator_test.py @@ -9,7 +9,7 @@ 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, @@ -436,6 +436,40 @@ def write_snapshot(gcs_path, content): self.assertEqual(materialized_task_config, task_config) self.assertEqual(materialized_resource_config, resource_config) + @patch("gigl.src.validation_check.config_validator.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") + source_configs = { + task_uri.uri: "shared_config:\n is_graph_directed: true\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, + should_create_symlinks_if_possible, + ): + 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() From 733f675102ab3a1cb129b0b5329b2262f05cb062 Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Thu, 30 Jul 2026 18:52:13 +0000 Subject: [PATCH 7/9] Move URI staging into config composition --- .../config_guides/hydra_composition.md | 16 ++++---- gigl/common/utils/proto_utils.py | 21 +++++++++-- gigl/src/validation_check/config_validator.py | 37 ++----------------- .../src/validation/config_validator_test.py | 2 +- 4 files changed, 31 insertions(+), 45 deletions(-) diff --git a/docs/user_guide/config_guides/hydra_composition.md b/docs/user_guide/config_guides/hydra_composition.md index e7175e905..0c4fd813e 100644 --- a/docs/user_guide/config_guides/hydra_composition.md +++ b/docs/user_guide/config_guides/hydra_composition.md @@ -3,13 +3,13 @@ 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 explicit local-only composition API. +is the URI-agnostic composition API. ## Config root -The primary file's parent directory is its Hydra config root. Composed primary files and selected fragments must use -Hydra's `.yaml` extension. Plain legacy `.yml` files remain supported by the non-composing reader, but ConfigValidator -inputs must use `.yaml`. +The primary file's parent directory is its Hydra config root. Local composed primary files and selected fragments must +use Hydra's `.yaml` extension. Plain legacy `.yml` files remain supported by the non-composing reader. Remote primaries +are staged with a `.yaml` extension before composition. ```text configs/ @@ -22,9 +22,9 @@ configs/ ``` Keep primary files directly in the bundle root. GiGL does not search for a repository root or a directory named -`configs`. Defaults List composition currently requires a local config bundle. ConfigValidator downloads GCS and HTTP -primaries to a temporary local file before composing them, so remote primaries remain supported only as plain, -single-file YAML. +`configs`. Local configs compose in place 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, so remote +primaries support only single-file composition. ## Task config example @@ -95,5 +95,5 @@ and runner intentionally differ. ## Boundaries - GiGL does not consume Hydra command-line overrides, multirun, launchers, or output-directory behavior. -- GCS and HTTP configs remain supported only when they are plain, single-file YAML. +- GCS and HTTP configs support single-file composition, but cannot select sibling config fragments. - Treat config bundle write access as trusted access. GiGL configs can reference importable classes and commands. diff --git a/gigl/common/utils/proto_utils.py b/gigl/common/utils/proto_utils.py index dc144cc14..2f53f5afc 100644 --- a/gigl/common/utils/proto_utils.py +++ b/gigl/common/utils/proto_utils.py @@ -50,9 +50,24 @@ 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: LocalUri, proto_cls: Type[T]) -> T: - """Compose a local YAML config with Hydra and parse it as a protobuf.""" - obj_dict = compose_yaml_config(uri=uri) + 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): + local_uri = uri + obj_dict = compose_yaml_config(uri=local_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, + should_create_symlinks_if_possible=False, + ) + obj_dict = compose_yaml_config(uri=local_uri) proto = ParseDict(js_dict=cast(dict, obj_dict), message=proto_cls()) return proto diff --git a/gigl/src/validation_check/config_validator.py b/gigl/src/validation_check/config_validator.py index 8d7759185..78509452b 100644 --- a/gigl/src/validation_check/config_validator.py +++ b/gigl/src/validation_check/config_validator.py @@ -1,11 +1,8 @@ import argparse from pathlib import Path -from tempfile import NamedTemporaryFile -from typing import Optional, Type, TypeVar +from typing import Optional -from google.protobuf import message - -from gigl.common import GcsUri, LocalUri, Uri, UriFactory +from gigl.common import GcsUri, Uri, UriFactory from gigl.common.logger import Logger from gigl.common.utils.gcs import GcsUtils from gigl.common.utils.proto_utils import ProtoUtils, proto_to_yaml @@ -14,7 +11,6 @@ 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.gigl_runtime import initialize_gigl_runtime from gigl.src.validation_check.libs.frozen_config_path_checks import ( assert_preprocessed_metadata_exists, @@ -202,8 +198,6 @@ logger = Logger() -T = TypeVar("T", bound=message.Message) - # 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 @@ -368,13 +362,11 @@ def resolve_configs( ) -> tuple[gbml_config_pb2.GbmlConfig, GiglResourceConfig]: """Resolve task and resource configs into self-contained protobufs.""" proto_utils = ProtoUtils() - task_config = _compose_proto_from_uri( - proto_utils=proto_utils, + task_config = proto_utils.compose_proto_from_yaml( uri=task_config_uri, proto_cls=gbml_config_pb2.GbmlConfig, ) - resource_config = _compose_proto_from_uri( - proto_utils=proto_utils, + resource_config = proto_utils.compose_proto_from_yaml( uri=resource_config_uri, proto_cls=GiglResourceConfig, ) @@ -384,27 +376,6 @@ def resolve_configs( return task_config, resource_config -def _compose_proto_from_uri( - proto_utils: ProtoUtils, - uri: Uri, - proto_cls: Type[T], -) -> T: - if isinstance(uri, LocalUri): - return proto_utils.compose_proto_from_yaml(uri=uri, proto_cls=proto_cls) - - with NamedTemporaryFile(suffix=".yaml") as temp_file: - local_uri = LocalUri(temp_file.name) - FileLoader().load_file( - file_uri_src=uri, - file_uri_dst=local_uri, - should_create_symlinks_if_possible=False, - ) - return proto_utils.compose_proto_from_yaml( - uri=local_uri, - proto_cls=proto_cls, - ) - - def kfp_validation_checks( job_name: str, task_config_uri: Uri, diff --git a/tests/unit/src/validation/config_validator_test.py b/tests/unit/src/validation/config_validator_test.py index 992760b8a..3bad34343 100644 --- a/tests/unit/src/validation/config_validator_test.py +++ b/tests/unit/src/validation/config_validator_test.py @@ -436,7 +436,7 @@ def write_snapshot(gcs_path, content): self.assertEqual(materialized_task_config, task_config) self.assertEqual(materialized_resource_config, resource_config) - @patch("gigl.src.validation_check.config_validator.FileLoader.load_file") + @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") From df3f6cec608ebde2e5c1c8ff0e55f334bec53c5a Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Thu, 30 Jul 2026 20:28:16 +0000 Subject: [PATCH 8/9] Simplify config composition staging --- .../config_guides/hydra_composition.md | 12 +++--- gigl/common/utils/hydra_config.py | 39 +++++++++++-------- gigl/common/utils/proto_utils.py | 17 +++----- .../kubeflow/kfp_orchestrator.py | 16 +++----- gigl/orchestration/kubeflow/kfp_pipeline.py | 2 + gigl/src/validation_check/config_validator.py | 27 ++++++------- tests/unit/common/utils/proto_utils_test.py | 11 ------ .../src/validation/config_validator_test.py | 1 - 8 files changed, 54 insertions(+), 71 deletions(-) diff --git a/docs/user_guide/config_guides/hydra_composition.md b/docs/user_guide/config_guides/hydra_composition.md index 0c4fd813e..03b4de5d0 100644 --- a/docs/user_guide/config_guides/hydra_composition.md +++ b/docs/user_guide/config_guides/hydra_composition.md @@ -7,9 +7,9 @@ is the URI-agnostic composition API. ## Config root -The primary file's parent directory is its Hydra config root. Local composed primary files and selected fragments must -use Hydra's `.yaml` extension. Plain legacy `.yml` files remain supported by the non-composing reader. Remote primaries -are staged with a `.yaml` extension before composition. +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/ @@ -22,9 +22,9 @@ configs/ ``` Keep primary files directly in the bundle root. GiGL does not search for a repository root or a directory named -`configs`. Local configs compose in place 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, so remote -primaries support only single-file composition. +`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, so remote primaries support only single-file composition. ## Task config example diff --git a/gigl/common/utils/hydra_config.py b/gigl/common/utils/hydra_config.py index 44213a353..c64607ce4 100644 --- a/gigl/common/utils/hydra_config.py +++ b/gigl/common/utils/hydra_config.py @@ -17,14 +17,6 @@ _COMPOSE_LOCK = threading.RLock() -def _register_gigl_now_resolver() -> None: - OmegaConf.register_new_resolver( - "now", - now_resolver, - replace=True, - ) - - def compose_yaml_config(uri: LocalUri) -> dict[str, Any]: """Compose a YAML config with Hydra using its parent as the config root. @@ -35,16 +27,23 @@ def compose_yaml_config(uri: LocalUri) -> dict[str, Any]: A fully composed and resolved mapping. Raises: - ValueError: If the primary filename is unsupported or the result is not - a mapping. + ValueError: If the result is not a mapping. RuntimeError: If another Hydra application owns the global context. """ - primary_name = uri.get_basename() - if not primary_name.endswith(".yaml"): - raise ValueError(f"Hydra primary config must use the .yaml extension: {uri}") + primary_path = Path(uri.uri).absolute() + if primary_path.is_symlink(): + # Local staging creates this symlink. Follow only that hop so a source + # URI that is itself a symlink keeps its filename and config root. + symlink_target = primary_path.readlink() + primary_path = ( + symlink_target + if symlink_target.is_absolute() + else primary_path.parent / symlink_target + ) + primary_name = primary_path.name config_name = primary_name.rsplit(".", 1)[0] - config_root = Path(uri.uri).absolute().parent + config_root = primary_path.parent with _COMPOSE_LOCK: if GlobalHydra.instance().is_initialized(): @@ -61,11 +60,19 @@ def compose_yaml_config(uri: LocalUri) -> dict[str, Any]: # initialization. GiGL's resolver is a backward-compatible # superset that also supports offsets. register_resolvers() - _register_gigl_now_resolver() + OmegaConf.register_new_resolver( + "now", + now_resolver, + replace=True, + ) composed = compose(config_name=config_name, overrides=[]) resolved = OmegaConf.to_container(composed, resolve=True) finally: - _register_gigl_now_resolver() + OmegaConf.register_new_resolver( + "now", + now_resolver, + replace=True, + ) if not isinstance(resolved, dict): raise ValueError( diff --git a/gigl/common/utils/proto_utils.py b/gigl/common/utils/proto_utils.py index 2f53f5afc..a2e7393b2 100644 --- a/gigl/common/utils/proto_utils.py +++ b/gigl/common/utils/proto_utils.py @@ -56,18 +56,13 @@ def compose_proto_from_yaml(self, uri: Uri, proto_cls: Type[T]) -> T: Remote URIs are downloaded as a single primary config; relative Defaults List entries are not fetched. """ - if isinstance(uri, LocalUri): - local_uri = uri + 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) - 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, - should_create_symlinks_if_possible=False, - ) - obj_dict = compose_yaml_config(uri=local_uri) proto = ParseDict(js_dict=cast(dict, obj_dict), message=proto_cls()) return proto diff --git a/gigl/orchestration/kubeflow/kfp_orchestrator.py b/gigl/orchestration/kubeflow/kfp_orchestrator.py index 1f7230a99..56638d2bd 100644 --- a/gigl/orchestration/kubeflow/kfp_orchestrator.py +++ b/gigl/orchestration/kubeflow/kfp_orchestrator.py @@ -13,7 +13,6 @@ from gigl.common.services.vertex_ai import VertexAIService from gigl.common.types.resource_config import CommonPipelineComponentConfigs from gigl.common.utils.proto_utils import ProtoUtils -from gigl.env.pipelines_config import get_resource_config from gigl.orchestration.kubeflow.kfp_pipeline import generate_pipeline from gigl.src.common.constants.components import GiGLComponents from gigl.src.common.types import AppliedTaskIdentifier @@ -148,17 +147,12 @@ def run( ) logger.info(f"Skipping pipeline compilation; will use {compiled_pipeline_path}") - if isinstance(resource_config_uri, LocalUri): - resource_config = GiglResourceConfigWrapper( - ProtoUtils().compose_proto_from_yaml( - uri=resource_config_uri, - proto_cls=GiglResourceConfig, - ) - ) - else: - resource_config = get_resource_config( - resource_config_uri=resource_config_uri + 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, diff --git a/gigl/orchestration/kubeflow/kfp_pipeline.py b/gigl/orchestration/kubeflow/kfp_pipeline.py index 8ba3659c1..6fb3f3dce 100644 --- a/gigl/orchestration/kubeflow/kfp_pipeline.py +++ b/gigl/orchestration/kubeflow/kfp_pipeline.py @@ -157,6 +157,8 @@ def _generate_component_tasks( resource_config_uri=resource_config_uri, common_pipeline_component_configs=common_pipeline_component_configs, ) + # 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[ diff --git a/gigl/src/validation_check/config_validator.py b/gigl/src/validation_check/config_validator.py index 78509452b..0f8c8fdfc 100644 --- a/gigl/src/validation_check/config_validator.py +++ b/gigl/src/validation_check/config_validator.py @@ -370,9 +370,12 @@ def resolve_configs( uri=resource_config_uri, proto_cls=GiglResourceConfig, ) - resource_config.shared_resource_config.CopyFrom( - GiglResourceConfigWrapper(resource_config).shared_resource_config - ) + 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 @@ -413,19 +416,13 @@ def materialize_resolved_configs( The resolved task and resource config URIs. """ resource_config_wrapper = GiglResourceConfigWrapper(resource_config) - snapshot_root = GcsUri.join( - resource_config_wrapper.temp_assets_regional_bucket_path, - job_name, - "config_validator", - ) - task_config_uri = GcsUri.join( - snapshot_root, - "resolved_task_config.yaml", - ) - resource_config_uri = GcsUri.join( - snapshot_root, - "resolved_resource_config.yaml", + 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, diff --git a/tests/unit/common/utils/proto_utils_test.py b/tests/unit/common/utils/proto_utils_test.py index eaa4ae67a..126e5b29a 100644 --- a/tests/unit/common/utils/proto_utils_test.py +++ b/tests/unit/common/utils/proto_utils_test.py @@ -291,17 +291,6 @@ def test_resource_composition_resolves_dynamic_nested_default(self): "example-project", ) - def test_composed_primary_requires_yaml_extension(self): - with TemporaryDirectory() as temp_directory: - config_path = Path(temp_directory) / "task.yml" - config_path.write_text("defaults:\n - _self_\n") - - with self.assertRaises(ValueError): - self.proto_utils.compose_proto_from_yaml( - uri=LocalUri(config_path), - proto_cls=gbml_config_pb2.GbmlConfig, - ) - def test_composition_is_thread_safe(self): with TemporaryDirectory() as temp_directory: config_root = Path(temp_directory) diff --git a/tests/unit/src/validation/config_validator_test.py b/tests/unit/src/validation/config_validator_test.py index 3bad34343..925de4c96 100644 --- a/tests/unit/src/validation/config_validator_test.py +++ b/tests/unit/src/validation/config_validator_test.py @@ -453,7 +453,6 @@ def test_downloads_remote_configs_before_composing(self, mock_load_file) -> None def download_config( file_uri_src, file_uri_dst, - should_create_symlinks_if_possible, ): self.assertTrue(file_uri_dst.get_basename().endswith(".yaml")) Path(file_uri_dst.uri).write_text(source_configs[file_uri_src.uri]) From f29978c63317fa58e9a0e8d55d2f8744064166b6 Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Thu, 30 Jul 2026 20:37:00 +0000 Subject: [PATCH 9/9] Document remote Hydra search paths --- .../config_guides/hydra_composition.md | 25 +++++++++++++++++-- gigl/common/utils/hydra_config.py | 14 ++--------- gigl/common/utils/proto_utils.py | 19 ++++++++------ .../src/validation/config_validator_test.py | 14 ++++++++++- 4 files changed, 50 insertions(+), 22 deletions(-) diff --git a/docs/user_guide/config_guides/hydra_composition.md b/docs/user_guide/config_guides/hydra_composition.md index 03b4de5d0..9ef4002d4 100644 --- a/docs/user_guide/config_guides/hydra_composition.md +++ b/docs/user_guide/config_guides/hydra_composition.md @@ -24,7 +24,28 @@ configs/ 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, so remote primaries support only single-file composition. +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 @@ -95,5 +116,5 @@ and runner intentionally differ. ## Boundaries - GiGL does not consume Hydra command-line overrides, multirun, launchers, or output-directory behavior. -- GCS and HTTP configs support single-file composition, but cannot select sibling config fragments. +- 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/gigl/common/utils/hydra_config.py b/gigl/common/utils/hydra_config.py index c64607ce4..d9ef817f7 100644 --- a/gigl/common/utils/hydra_config.py +++ b/gigl/common/utils/hydra_config.py @@ -30,20 +30,10 @@ def compose_yaml_config(uri: LocalUri) -> dict[str, Any]: ValueError: If the result is not a mapping. RuntimeError: If another Hydra application owns the global context. """ - primary_path = Path(uri.uri).absolute() - if primary_path.is_symlink(): - # Local staging creates this symlink. Follow only that hop so a source - # URI that is itself a symlink keeps its filename and config root. - symlink_target = primary_path.readlink() - primary_path = ( - symlink_target - if symlink_target.is_absolute() - else primary_path.parent / symlink_target - ) - primary_name = primary_path.name + primary_name = uri.get_basename() config_name = primary_name.rsplit(".", 1)[0] - config_root = primary_path.parent + config_root = Path(uri.uri).absolute().parent with _COMPOSE_LOCK: if GlobalHydra.instance().is_initialized(): diff --git a/gigl/common/utils/proto_utils.py b/gigl/common/utils/proto_utils.py index a2e7393b2..4f3b2fd33 100644 --- a/gigl/common/utils/proto_utils.py +++ b/gigl/common/utils/proto_utils.py @@ -56,13 +56,18 @@ def compose_proto_from_yaml(self, uri: Uri, proto_cls: Type[T]) -> T: Remote URIs are downloaded as a single primary config; relative Defaults List entries are not fetched. """ - 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) + 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 diff --git a/tests/unit/src/validation/config_validator_test.py b/tests/unit/src/validation/config_validator_test.py index 925de4c96..71e7d5058 100644 --- a/tests/unit/src/validation/config_validator_test.py +++ b/tests/unit/src/validation/config_validator_test.py @@ -440,8 +440,20 @@ def write_snapshot(gcs_path, content): 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: "shared_config:\n is_graph_directed: true\n", + 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"