From af93afb5eabbf93ea498b38e1db62ad9c2613cc1 Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Thu, 3 Sep 2026 23:10:54 +0000 Subject: [PATCH 1/5] Prepare for Python 3.12 and 3.13: drop distutils, make GLT install failures fatal Python 3.12 removed `distutils` and the deprecated `unittest` aliases, and GiGL uses both. The `distutils` imports work today only because `setuptools` is installed transitively and ships a compatibility shim; a library must not depend on that. - Add `gigl.common.utils.parse.str_to_bool` and use it in place of `distutils.util.strtobool` at all 15 call sites. Accepted and rejected spellings match `strtobool`. Every call site already coerced the result with `bool()` or used it as a condition, so the `int` to `bool` return change is not observable. - Rename the 26 `assertEquals` / `assertNotEquals` uses to `assertEqual` / `assertNotEqual`, and select ruff `UP005` so they cannot return. - Make a failed `install_glt.sh` fatal. `main()` returned the child's status but the `__main__` block discarded it, so `requirements/install_py_deps.sh` saw exit 0 under `set -e` and every base image build continued after a failed GLT install. Measured against a stub that exits 7: the old script exits 0, the new one exits 7. A successful install still exits 0. - Run `ty` twice in `make type_check`, at the 3.11 floor and at 3.13. `ty` resolves the standard library against one version per invocation, so the floor pass accepts modules 3.13 removed and only the ceiling pass rejects them. Co-Authored-By: Claude Opus 5 (1M context) --- Makefile | 6 ++ .../graph_store/storage_main.py | 4 +- .../KDD_2025/heterogeneous_inference.py | 4 +- .../KDD_2025/heterogeneous_training.py | 4 +- gigl/common/utils/parse.py | 41 ++++++++++++ gigl/distributed/dataset_factory.py | 6 +- gigl/scripts/post_install.py | 28 ++++++-- ...ased_link_prediction_modeling_task_spec.py | 4 +- .../utils/profiler_wrapper.py | 8 +-- .../common/types/pb_wrappers/gbml_config.py | 8 +-- gigl/src/subgraph_sampler/subgraph_sampler.py | 4 +- gigl/src/training/v1/lib/training_process.py | 6 +- pyproject.toml | 10 ++- scripts/bootstrap_resource_config.py | 4 +- tests/integration/common/dataflow_test.py | 8 +-- tests/integration/common/gcs_test.py | 2 +- .../data_preprocessor_pipeline_test.py | 8 +-- .../pipeline/inferencer/inferencer_test.py | 4 +- .../split_generator_pipeline_test.py | 2 +- .../subgraph_sampler/subgraph_sampler_test.py | 2 +- .../unit/common/collections/itertools_test.py | 2 +- tests/unit/common/utils/parse_test.py | 50 ++++++++++++++ tests/unit/common/utils/retry_test.py | 4 +- tests/unit/scripts/__init__.py | 0 tests/unit/scripts/post_install_test.py | 67 +++++++++++++++++++ .../graph_builder/pyg_graph_builder_test.py | 10 +-- .../graph_builder/pyg_graph_data_test.py | 8 +-- .../tf_records_iterable_dataset_test.py | 2 +- 28 files changed, 248 insertions(+), 58 deletions(-) create mode 100644 gigl/common/utils/parse.py create mode 100644 tests/unit/common/utils/parse_test.py create mode 100644 tests/unit/scripts/__init__.py create mode 100644 tests/unit/scripts/post_install_test.py diff --git a/Makefile b/Makefile index 48211baec..1e0609c9d 100644 --- a/Makefile +++ b/Makefile @@ -173,8 +173,14 @@ format_cpp: format: format_py format_cpp format_scala format_md +# ty resolves the stdlib against a single Python version per invocation, so each end of +# the supported range needs its own pass. type_check: +# Floor: 3.11, from [tool.ty.environment] in pyproject.toml. Catches stdlib APIs that do +# not exist that far back. uv run ty check ${PYTHON_DIRS} +# Ceiling: catches stdlib modules and signatures 3.13 removed or changed. + uv run ty check --python-version 3.13 ${PYTHON_DIRS} build_cpp_extensions: $(MAKE) -C gigl-core build_cpp_extensions diff --git a/examples/link_prediction/graph_store/storage_main.py b/examples/link_prediction/graph_store/storage_main.py index 80ed3aed0..8c55759ef 100644 --- a/examples/link_prediction/graph_store/storage_main.py +++ b/examples/link_prediction/graph_store/storage_main.py @@ -74,7 +74,6 @@ import argparse import ast import os -from distutils.util import strtobool from typing import Literal, Optional, Union import torch @@ -82,6 +81,7 @@ from gigl.common import Uri, UriFactory from gigl.common.logger import Logger from gigl.common.utils.os_utils import import_obj +from gigl.common.utils.parse import str_to_bool from gigl.distributed.graph_store import ( GraphStoreInfo, build_storage_dataset, @@ -254,7 +254,7 @@ def storage_node_process( splitter=splitter, ssl_positive_label_percentage=ssl_positive_label_percentage, should_load_tf_records_in_parallel=bool( - strtobool(args.should_load_tf_records_in_parallel) + str_to_bool(args.should_load_tf_records_in_parallel) ), num_rpc_threads=args.num_rpc_threads, rpc_timeout=args.rpc_timeout, diff --git a/examples/tutorial/KDD_2025/heterogeneous_inference.py b/examples/tutorial/KDD_2025/heterogeneous_inference.py index 772d695e7..e4b1a5a69 100644 --- a/examples/tutorial/KDD_2025/heterogeneous_inference.py +++ b/examples/tutorial/KDD_2025/heterogeneous_inference.py @@ -26,7 +26,6 @@ import argparse import datetime from collections.abc import Mapping -from distutils.util import strtobool from pathlib import Path import fastavro @@ -38,6 +37,7 @@ from gigl.common import Uri, UriFactory from gigl.common.data.export import EmbeddingExporter from gigl.common.logger import Logger +from gigl.common.utils.parse import str_to_bool from gigl.distributed import ( DistDataset, DistNeighborLoader, @@ -154,7 +154,7 @@ def inference( task_config_uri, _tfrecord_uri_pattern=".*tfrecord", ) - if strtobool(args.use_local_saved_model): + if str_to_bool(args.use_local_saved_model): model_uri = LOCAL_SAVED_MODEL_URI else: model_uri = gbml_config_pb_wrapper.shared_config.trained_model_metadata.trained_model_uri diff --git a/examples/tutorial/KDD_2025/heterogeneous_training.py b/examples/tutorial/KDD_2025/heterogeneous_training.py index 7035fff32..c6ead56dd 100644 --- a/examples/tutorial/KDD_2025/heterogeneous_training.py +++ b/examples/tutorial/KDD_2025/heterogeneous_training.py @@ -37,7 +37,6 @@ import argparse from collections.abc import Iterable, Mapping -from distutils.util import strtobool from typing import Literal import torch @@ -48,6 +47,7 @@ from examples.tutorial.KDD_2025.utils import LOCAL_SAVED_MODEL_URI, init_model from gigl.common import UriFactory from gigl.common.logger import Logger +from gigl.common.utils.parse import str_to_bool from gigl.distributed import ( DistABLPLoader, DistDataset, @@ -269,7 +269,7 @@ def train( logger.info(f"Test node type {node_type} has {node_ids.size(0)} nodes.") # ty: ignore[unresolved-attribute] TODO(ty-torch-keyed-access): fix ty false positives for torch-backed keyed container access. training_process_port = get_free_port() logger.info(f"Will train for {max_training_batches} batches.") - if strtobool(args.use_local_saved_model): + if str_to_bool(args.use_local_saved_model): model_uri = LOCAL_SAVED_MODEL_URI else: model_uri = gbml_config_pb_wrapper.shared_config.trained_model_metadata.trained_model_uri diff --git a/gigl/common/utils/parse.py b/gigl/common/utils/parse.py new file mode 100644 index 000000000..a6de44f8d --- /dev/null +++ b/gigl/common/utils/parse.py @@ -0,0 +1,41 @@ +from typing import Final + +_TRUE_VALUES: Final[frozenset[str]] = frozenset({"y", "yes", "t", "true", "on", "1"}) +_FALSE_VALUES: Final[frozenset[str]] = frozenset({"n", "no", "f", "false", "off", "0"}) + + +def str_to_bool(value: str) -> bool: + """ + Converts a string representation of truth to a bool. + + Accepts and rejects the same spellings as `distutils.util.strtobool`, which is no + longer part of the standard library, though the `ValueError` message keeps the + caller's casing instead of lowercasing it: + + - True: "y", "yes", "t", "true", "on", "1" + - False: "n", "no", "f", "false", "off", "0" + + Matching is case-insensitive. Whitespace is not stripped, so " true" raises rather + than returning True. + + Example: + >>> str_to_bool("TRUE") + True + >>> str_to_bool("off") + False + + Args: + value (str): The string to interpret as a boolean. + + Returns: + bool: The truth value `value` spells. + + Raises: + ValueError: If `value` is not one of the accepted spellings. + """ + normalized_value = value.lower() + if normalized_value in _TRUE_VALUES: + return True + if normalized_value in _FALSE_VALUES: + return False + raise ValueError(f"invalid truth value {value!r}") diff --git a/gigl/distributed/dataset_factory.py b/gigl/distributed/dataset_factory.py index 0ffa3c462..6cc131596 100644 --- a/gigl/distributed/dataset_factory.py +++ b/gigl/distributed/dataset_factory.py @@ -5,7 +5,6 @@ import time from collections.abc import Mapping -from distutils.util import strtobool from typing import Literal, MutableMapping, Optional, Tuple, Type, Union import torch @@ -29,6 +28,7 @@ ) from gigl.common.logger import Logger from gigl.common.utils.decorator import tf_on_cpu +from gigl.common.utils.parse import str_to_bool from gigl.distributed.constants import DEFAULT_MASTER_DATA_BUILDING_PORT from gigl.distributed.dist_context import DistributedContext from gigl.distributed.dist_dataset import DistDataset @@ -629,11 +629,11 @@ def build_dataset_from_task_config_uri( ) should_use_range_partitioning = bool( - strtobool(args.get("should_use_range_partitioning", "True")) + str_to_bool(args.get("should_use_range_partitioning", "True")) ) should_load_tensors_in_parallel = bool( - strtobool(args.get("should_load_tensors_in_parallel", "True")) + str_to_bool(args.get("should_load_tensors_in_parallel", "True")) ) logger.info( diff --git a/gigl/scripts/post_install.py b/gigl/scripts/post_install.py index cd13c4b0b..c1974a61a 100644 --- a/gigl/scripts/post_install.py +++ b/gigl/scripts/post_install.py @@ -34,8 +34,17 @@ def run_command_and_stream_stdout(cmd: str) -> Optional[int]: return return_code -def main(): - """Main entry point for the post-install script.""" +def main() -> int: + """Main entry point for the post-install script. + + Returns: + int: 0, when install_glt.sh succeeds. + + Raises: + SystemExit: With a non-zero code when install_glt.sh is missing, fails, or + reports no exit status. Callers must propagate it: swallowing it lets a + build succeed while shipping an environment with no working GLT. + """ print("Running GIGL post-install script...") # Get the directory where this script is located @@ -52,9 +61,16 @@ def main(): try: print(f"Executing {cmd}...") - result = run_command_and_stream_stdout(cmd) - print("Post-install script finished running, with return code: ", result) - return result + return_code = run_command_and_stream_stdout(cmd) + print("Post-install script finished running, with return code: ", return_code) + # `Popen.poll()` returns None while the child has no recorded status, so an + # unknown outcome is not evidence of success. + if return_code is None: + print("Error: could not determine the exit status of install_glt.sh") + sys.exit(1) + if return_code != 0: + sys.exit(return_code) + return return_code except subprocess.CalledProcessError as e: print(f"Error running install_glt.sh: {e}") @@ -65,4 +81,4 @@ def main(): if __name__ == "__main__": - main() + raise SystemExit(main()) diff --git a/gigl/src/common/modeling_task_specs/node_anchor_based_link_prediction_modeling_task_spec.py b/gigl/src/common/modeling_task_specs/node_anchor_based_link_prediction_modeling_task_spec.py index c8fd60e21..80227b31c 100644 --- a/gigl/src/common/modeling_task_specs/node_anchor_based_link_prediction_modeling_task_spec.py +++ b/gigl/src/common/modeling_task_specs/node_anchor_based_link_prediction_modeling_task_spec.py @@ -1,6 +1,5 @@ from collections import OrderedDict from contextlib import ExitStack -from distutils.util import strtobool from time import time from typing import Any, Optional, Type @@ -12,6 +11,7 @@ from gigl.common.logger import Logger from gigl.common.utils import os_utils +from gigl.common.utils.parse import str_to_bool from gigl.common.utils.torch_training import ( get_rank, get_world_size, @@ -201,7 +201,7 @@ def __init__(self, **kwargs) -> None: # Retrieval-specific Task Parameters softmax_temp = float(kwargs.get("softmax_temp", 0.07)) should_remove_accidental_hits = bool( - strtobool(kwargs.get("should_remove_accidental_hits", "True")) + str_to_bool(kwargs.get("should_remove_accidental_hits", "True")) ) task = base_task( temperature=softmax_temp, diff --git a/gigl/src/common/modeling_task_specs/utils/profiler_wrapper.py b/gigl/src/common/modeling_task_specs/utils/profiler_wrapper.py index 2a7c423a5..04f28cdec 100644 --- a/gigl/src/common/modeling_task_specs/utils/profiler_wrapper.py +++ b/gigl/src/common/modeling_task_specs/utils/profiler_wrapper.py @@ -1,5 +1,4 @@ import tempfile -from distutils.util import strtobool from torch.profiler import ( ProfilerActivity, @@ -10,6 +9,7 @@ from gigl.common import LocalUri from gigl.common.logger import Logger +from gigl.common.utils.parse import str_to_bool logger = Logger() @@ -33,9 +33,9 @@ def __init__(self, **kwargs) -> None: active=self.active, repeat=self.repeat, ) - self.profile_memory = bool(strtobool(kwargs.get("profile_memory", "True"))) - self.record_shapes = bool(strtobool(kwargs.get("record_shapes", "False"))) - self.with_stack = bool(strtobool(kwargs.get("with_stack", "False"))) + self.profile_memory = bool(str_to_bool(kwargs.get("profile_memory", "True"))) + self.record_shapes = bool(str_to_bool(kwargs.get("record_shapes", "False"))) + self.with_stack = bool(str_to_bool(kwargs.get("with_stack", "False"))) logger.info(f"Profiler will be instantiated with {self.__dict__}") def profiler_context(self) -> profile: diff --git a/gigl/src/common/types/pb_wrappers/gbml_config.py b/gigl/src/common/types/pb_wrappers/gbml_config.py index f153a6805..00915267e 100644 --- a/gigl/src/common/types/pb_wrappers/gbml_config.py +++ b/gigl/src/common/types/pb_wrappers/gbml_config.py @@ -1,11 +1,11 @@ from __future__ import annotations from dataclasses import dataclass, field -from distutils.util import strtobool from typing import Optional from gigl.common import Uri, UriFactory from gigl.common.logger import Logger +from gigl.common.utils.parse import str_to_bool from gigl.common.utils.proto_utils import ProtoUtils from gigl.src.common.types.graph_data import EdgeType, NodeType from gigl.src.common.types.pb_wrappers.dataset_metadata import DatasetMetadataPbWrapper @@ -529,7 +529,7 @@ def should_use_glt_backend(self) -> bool: """ return bool( - strtobool( + str_to_bool( dict(self.gbml_config_pb.feature_flags).get( "should_run_glt_backend", "False" ) @@ -550,7 +550,7 @@ def should_populate_predictions_path(self) -> bool: bool: Whether to populate predictions path in the InferenceOutput for each entity type """ return bool( - strtobool( + str_to_bool( dict(self.gbml_config_pb.feature_flags).get( "should_populate_predictions_path", "False" ) @@ -570,7 +570,7 @@ def should_populate_embeddings_path(self) -> bool: bool: Whether to populate embeddings path in the InferenceOutput for each entity type """ return bool( - strtobool( + str_to_bool( dict(self.gbml_config_pb.feature_flags).get( "should_populate_embeddings_path", "True" ) diff --git a/gigl/src/subgraph_sampler/subgraph_sampler.py b/gigl/src/subgraph_sampler/subgraph_sampler.py index 359418e8a..73f80c5a0 100644 --- a/gigl/src/subgraph_sampler/subgraph_sampler.py +++ b/gigl/src/subgraph_sampler/subgraph_sampler.py @@ -1,7 +1,6 @@ import argparse import datetime import os -from distutils.util import strtobool from typing import Optional, Sequence import gigl.env.dep_constants as dep_constants @@ -16,6 +15,7 @@ from gigl.common.metrics.decorators import flushes_metrics, profileit from gigl.common.utils import os_utils from gigl.common.utils.gcs import GcsUtils +from gigl.common.utils.parse import str_to_bool from gigl.env.pipelines_config import get_resource_config from gigl.src.common.constants.components import GiGLComponents from gigl.src.common.constants.metrics import TIMER_SUBGRAPH_SAMPLER_S @@ -101,7 +101,7 @@ def run( # Dataproc image 2.0 starting 2026-08-25. Setting the `use_spark35_runner` # experimental flag to "False" remains a temporary escape hatch until then. use_spark35: bool = bool( - strtobool( + str_to_bool( gbml_config_pb_wrapper.dataset_config.subgraph_sampler_config.experimental_flags.get( "use_spark35_runner", "True" ) diff --git a/gigl/src/training/v1/lib/training_process.py b/gigl/src/training/v1/lib/training_process.py index c79bd6983..393783fbe 100644 --- a/gigl/src/training/v1/lib/training_process.py +++ b/gigl/src/training/v1/lib/training_process.py @@ -5,7 +5,6 @@ import sys import tempfile import traceback -from distutils.util import strtobool from typing import Any, Optional import tensorflow as tf @@ -19,6 +18,7 @@ from gigl.common.metrics.decorators import flushes_metrics, profileit from gigl.common.utils import os_utils, torch_training from gigl.common.utils.local_fs import does_path_exist +from gigl.common.utils.parse import str_to_bool from gigl.common.utils.torch_training import ( get_distributed_backend, get_rank, @@ -300,7 +300,9 @@ def __run( # If all parameters are always expected to receive backprop in training, it is not recommended to enable this flag, as it can adversely affect # performance as a result of the extra traversal of the autograd graph every iteration. should_enable_find_unused_parameters = bool( - strtobool(trainer_args.get("should_enable_find_unused_parameters", "False")) + str_to_bool( + trainer_args.get("should_enable_find_unused_parameters", "False") + ) ) trainer_instance.model = setup_model_device( diff --git a/pyproject.toml b/pyproject.toml index 7fc6e79a7..974f409e8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -288,7 +288,13 @@ exclude = ["*_pb2.py", "*_pb2.pyi"] # Enforces a consistent import order (stdlib → third-party → first-party). # Replaces: isort # Docs: https://docs.astral.sh/ruff/rules/#isort-i -select = ["F401", "I"] +# +# UP005 - Deprecated unittest aliases (pyupgrade) +# Flags 15 of the 16 aliases removed in Python 3.12, such as `assertEquals`; it does +# not cover `assertDictContainsSubset`. Autofixable. +# Replaces: nothing. +# Docs: https://docs.astral.sh/ruff/rules/deprecated-unittest-alias/ +select = ["F401", "I", "UP005"] [tool.ruff.lint.per-file-ignores] # __init__.py files re-export symbols for the public API, so unused-import @@ -300,6 +306,8 @@ select = ["F401", "I"] known-first-party = ["gigl", "tests", "snapchat", "scripts"] [tool.ty.environment] +# The floor of the supported range. The `type_check` Make target covers the ceiling with +# a second `ty check --python-version 3.13` pass. python-version = "3.11" [tool.ty.src] diff --git a/scripts/bootstrap_resource_config.py b/scripts/bootstrap_resource_config.py index f3517138e..4bc8f06ff 100644 --- a/scripts/bootstrap_resource_config.py +++ b/scripts/bootstrap_resource_config.py @@ -6,12 +6,12 @@ import subprocess import tempfile from dataclasses import dataclass -from distutils.util import strtobool from typing import Optional import yaml from gigl.common import GcsUri, HttpUri, LocalUri, UriFactory +from gigl.common.utils.parse import str_to_bool from gigl.src.common.utils.file_loader import FileLoader GIGL_ROOT_DIR = pathlib.Path(__file__).resolve().parent.parent @@ -335,7 +335,7 @@ def assert_gcs_bucket_exists(bucket_name: str): print(f"Updated YAML file saved at '{destination_file_path}'") # Update the user's shell configuration - if args.force_shell_config_update and strtobool(args.force_shell_config_update): + if args.force_shell_config_update and str_to_bool(args.force_shell_config_update): should_update_shell_config = "y" print("Forcing shell updated due to --force_shell_config_update flag.") else: diff --git a/tests/integration/common/dataflow_test.py b/tests/integration/common/dataflow_test.py index 84b948789..d56c28c4b 100644 --- a/tests/integration/common/dataflow_test.py +++ b/tests/integration/common/dataflow_test.py @@ -50,7 +50,7 @@ def test_can_create_pipeline_config(self): # Ensure the pipeline options were propogated through parsed_options = options.get_all_options() - self.assertEquals(parsed_options["num_workers"], NUM_WORKERS) - self.assertEquals(parsed_options["max_num_workers"], MAX_NUM_WORKERS) - self.assertEquals(parsed_options["machine_type"], MACHINE_TYPE) - self.assertEquals(parsed_options["disk_size_gb"], DISK_SIZE_GB) + self.assertEqual(parsed_options["num_workers"], NUM_WORKERS) + self.assertEqual(parsed_options["max_num_workers"], MAX_NUM_WORKERS) + self.assertEqual(parsed_options["machine_type"], MACHINE_TYPE) + self.assertEqual(parsed_options["disk_size_gb"], DISK_SIZE_GB) diff --git a/tests/integration/common/gcs_test.py b/tests/integration/common/gcs_test.py index 8c2b020fd..cb3d713e4 100644 --- a/tests/integration/common/gcs_test.py +++ b/tests/integration/common/gcs_test.py @@ -21,7 +21,7 @@ def tearDown(self): gcs_utils.delete_files_in_bucket_dir(self._scratch_gcs_path) def test_join_path(self): - self.assertEquals( + self.assertEqual( GcsUri.join("gs://bucket_name", "path", "file.txt"), GcsUri("gs://bucket_name/path/file.txt"), ) diff --git a/tests/integration/pipeline/data_preprocessor/data_preprocessor_pipeline_test.py b/tests/integration/pipeline/data_preprocessor/data_preprocessor_pipeline_test.py index 9a651ede6..123ebc4ee 100644 --- a/tests/integration/pipeline/data_preprocessor/data_preprocessor_pipeline_test.py +++ b/tests/integration/pipeline/data_preprocessor/data_preprocessor_pipeline_test.py @@ -168,21 +168,21 @@ def __assert_graph_metadata_reflects_mocked_dataset_info( gbml_config_pb_wrapper.graph_metadata_pb_wrapper.condensed_edge_type_to_edge_type_map ) - self.assertEquals( + self.assertEqual( len(condensed_node_type_to_node_type_map), len(mocked_dataset_info.node_types), ) - self.assertEquals( + self.assertEqual( len(condensed_edge_type_to_edge_type_map), len(mocked_dataset_info.edge_types), ) - self.assertEquals( + self.assertEqual( condensed_node_type_to_node_type_map[DEFAULT_CONDENSED_NODE_TYPE], mocked_dataset_info.default_node_type, ) - self.assertEquals( + self.assertEqual( condensed_edge_type_to_edge_type_map[DEFAULT_CONDENSED_EDGE_TYPE].relation, mocked_dataset_info.default_edge_type.relation, ) diff --git a/tests/integration/pipeline/inferencer/inferencer_test.py b/tests/integration/pipeline/inferencer/inferencer_test.py index 905d996cc..b213c7573 100644 --- a/tests/integration/pipeline/inferencer/inferencer_test.py +++ b/tests/integration/pipeline/inferencer/inferencer_test.py @@ -199,7 +199,7 @@ def __validate_inferencer_for_mocked_dataset( node_type_to_inferencer_output_info_map[node_type].embeddings_path ) if should_assert_embeddings: - self.assertEquals( + self.assertEqual( self.__bq_utils.count_number_of_rows_in_bq_table( bq_table=node_type_to_inferencer_output_info_map[ node_type @@ -210,7 +210,7 @@ def __validate_inferencer_for_mocked_dataset( f"Found unexpected number of rows for node type {node_type} in embedding table.", ) if should_assert_predictions: - self.assertEquals( + self.assertEqual( self.__bq_utils.count_number_of_rows_in_bq_table( bq_table=node_type_to_inferencer_output_info_map[ node_type diff --git a/tests/integration/pipeline/split_generator/split_generator_pipeline_test.py b/tests/integration/pipeline/split_generator/split_generator_pipeline_test.py index 2a187992a..7401386d2 100644 --- a/tests/integration/pipeline/split_generator/split_generator_pipeline_test.py +++ b/tests/integration/pipeline/split_generator/split_generator_pipeline_test.py @@ -643,7 +643,7 @@ def __validate_node_classification_split( == supervised_node_classification.NodeClassificationSettingType.INDUCTIVE ): # All edge sets across train/val/test splits must be disjoint. - self.assertEquals( + self.assertEqual( train_graph.num_edges + val_graph.num_edges + test_graph.num_edges, composed_graph.num_edges, ) diff --git a/tests/integration/pipeline/subgraph_sampler/subgraph_sampler_test.py b/tests/integration/pipeline/subgraph_sampler/subgraph_sampler_test.py index 63ae84560..f4c439c2c 100644 --- a/tests/integration/pipeline/subgraph_sampler/subgraph_sampler_test.py +++ b/tests/integration/pipeline/subgraph_sampler/subgraph_sampler_test.py @@ -1337,7 +1337,7 @@ def __run_and_check_node_based_task_sgs_validity( supervision_node_type ] ) - self.assertEquals( + self.assertEqual( total_rooted_node_neighborhood_samples, expected_nodes_of_supervision_node_type, f"Found {total_rooted_node_neighborhood_samples} rooted samples from SGS output, but found {expected_nodes_of_supervision_node_type} nodes from Data Preprocessor output", diff --git a/tests/unit/common/collections/itertools_test.py b/tests/unit/common/collections/itertools_test.py index a494c6c36..ee68b12f8 100644 --- a/tests/unit/common/collections/itertools_test.py +++ b/tests/unit/common/collections/itertools_test.py @@ -7,4 +7,4 @@ def test_batch(self): input_list = [1, 2, 3, 4, 5] output = batch(list_of_items=input_list, chunk_size=2) expected_output = [[1, 2], [3, 4], [5]] - self.assertEquals(output, expected_output) + self.assertEqual(output, expected_output) diff --git a/tests/unit/common/utils/parse_test.py b/tests/unit/common/utils/parse_test.py new file mode 100644 index 000000000..1a4d9d51b --- /dev/null +++ b/tests/unit/common/utils/parse_test.py @@ -0,0 +1,50 @@ +import importlib + +from gigl.common.utils.parse import str_to_bool +from tests.test_assets.test_case import TestCase + +_TRUE_SPELLINGS = ("y", "yes", "t", "true", "on", "1") +_FALSE_SPELLINGS = ("n", "no", "f", "false", "off", "0") +# `" true"` belongs here because `str_to_bool` does not strip surrounding whitespace. +_INVALID_VALUES = ("", " true", "2", "none") + + +def _casings(spelling: str) -> tuple[str, ...]: + """All the casings `str_to_bool` must accept for one spelling.""" + return (spelling.lower(), spelling.upper(), spelling.capitalize()) + + +class ParseUtilsTest(TestCase): + def test_accepted_spellings(self) -> None: + for spelling in _TRUE_SPELLINGS: + for value in _casings(spelling): + with self.subTest(value=value): + self.assertIs(str_to_bool(value), True) + for spelling in _FALSE_SPELLINGS: + for value in _casings(spelling): + with self.subTest(value=value): + self.assertIs(str_to_bool(value), False) + + def test_rejects_unrecognized_values(self) -> None: + for value in _INVALID_VALUES: + with self.subTest(value=value): + self.assertRaises(ValueError, str_to_bool, value) + + def test_matches_distutils_strtobool(self) -> None: + # The import is dynamic because a static `from distutils.util import ...` fails the + # 3.13 pass of `make type_check`. The reference is whichever `distutils` is + # importable: CPython's on 3.11 without `setuptools`, otherwise the copy + # `setuptools` injects through `distutils-precedence.pth`, which is what GiGL's dev + # and build environments resolve on every Python version. + try: + strtobool = importlib.import_module("distutils.util").strtobool + except ImportError: + self.skipTest("distutils is not importable in this environment") + + for spelling in _TRUE_SPELLINGS + _FALSE_SPELLINGS: + for value in _casings(spelling): + with self.subTest(value=value): + self.assertEqual(str_to_bool(value), bool(strtobool(value))) + for value in _INVALID_VALUES: + with self.subTest(value=value): + self.assertRaises(ValueError, strtobool, value) diff --git a/tests/unit/common/utils/retry_test.py b/tests/unit/common/utils/retry_test.py index 7b4735065..d2f5e8409 100644 --- a/tests/unit/common/utils/retry_test.py +++ b/tests/unit/common/utils/retry_test.py @@ -37,7 +37,7 @@ def should_succeed_after_3_tries(): return True self.assertTrue(should_succeed_after_3_tries()) - self.assertEquals(exec_counter, 3) + self.assertEqual(exec_counter, 3) def test_retry_with_function_deadlines(self): exec_counter = 0 @@ -53,7 +53,7 @@ def should_timeout_first_try_and_then_succeed(): start = time() self.assertTrue(should_timeout_first_try_and_then_succeed()) - self.assertEquals(exec_counter, 2) + self.assertEqual(exec_counter, 2) total_time_s = time() - start self.assertLessEqual( total_time_s, 10 diff --git a/tests/unit/scripts/__init__.py b/tests/unit/scripts/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/unit/scripts/post_install_test.py b/tests/unit/scripts/post_install_test.py new file mode 100644 index 000000000..a7776ac3b --- /dev/null +++ b/tests/unit/scripts/post_install_test.py @@ -0,0 +1,67 @@ +"""Exit-code contract for `post_install.py` run as a script. + +This is the path `requirements/install_py_deps.sh` takes, so every Docker base image +build treats this exit code as the verdict on GLT: a zero exit publishes the image, and +a failed `install_glt.sh` that reports success ships an environment with no working GLT. +These tests run the real script as a subprocess against a stub `install_glt.sh` and +assert the process exit code, which is the only signal a build observes. The +`gigl-post-install` console script reaches `main()` by a different route and is not +covered here. +""" + +import shutil +import subprocess +import sys +import tempfile +from pathlib import Path +from typing import Optional + +import gigl.scripts.post_install +from tests.test_assets.test_case import TestCase + +# `post_install.py` resolves `install_glt.sh` as `Path(__file__).parent / "install_glt.sh"`, +# so a copy of the real script in a temp dir runs whichever stub sits beside it. +_POST_INSTALL_PATH: Path = Path(gigl.scripts.post_install.__file__) + + +class PostInstallTest(TestCase): + def _run_post_install( + self, install_glt_exit_code: Optional[int] + ) -> "subprocess.CompletedProcess[str]": + """Runs the real post-install script beside a stub `install_glt.sh`. + + Args: + install_glt_exit_code (Optional[int]): Status the stub exits with, or None to + leave the directory without an `install_glt.sh` at all. + + Returns: + subprocess.CompletedProcess[str]: The finished process, for its `returncode`. + """ + with tempfile.TemporaryDirectory() as script_dir: + script_path = Path(script_dir) / _POST_INSTALL_PATH.name + shutil.copyfile(_POST_INSTALL_PATH, script_path) + if install_glt_exit_code is not None: + # A two-line stub keeps the test hermetic: no network, no package install, + # no real GLT build. + (Path(script_dir) / "install_glt.sh").write_text( + f"#!/usr/bin/env bash\nexit {install_glt_exit_code}\n" + ) + return subprocess.run( + [sys.executable, str(script_path)], + capture_output=True, + text=True, + ) + + def test_exits_zero_when_install_glt_succeeds(self) -> None: + completed = self._run_post_install(install_glt_exit_code=0) + self.assertEqual(completed.returncode, 0, completed.stdout) + + def test_propagates_install_glt_exit_code(self) -> None: + # 7 is neither success nor the 1 the missing-script path uses, so matching it proves + # the child's own status reached the caller. + completed = self._run_post_install(install_glt_exit_code=7) + self.assertEqual(completed.returncode, 7, completed.stdout) + + def test_exits_one_when_install_glt_is_missing(self) -> None: + completed = self._run_post_install(install_glt_exit_code=None) + self.assertEqual(completed.returncode, 1, completed.stdout) diff --git a/tests/unit/src/common/graph_builder/pyg_graph_builder_test.py b/tests/unit/src/common/graph_builder/pyg_graph_builder_test.py index c164e3969..dbd0711c6 100644 --- a/tests/unit/src/common/graph_builder/pyg_graph_builder_test.py +++ b/tests/unit/src/common/graph_builder/pyg_graph_builder_test.py @@ -62,7 +62,7 @@ def test_can_create_accurate_graph_representation(self): ), } ) - self.assertEquals(graph_data_from_builder, expected_graph_data) + self.assertEqual(graph_data_from_builder, expected_graph_data) def test_can_create_with_with_no_edge_and_node_features(self): pyg_graph_builder = PygGraphBuilder() @@ -116,7 +116,7 @@ def test_can_create_with_with_no_edge_and_node_features(self): # This is a restriction of PyG, that is it expectes node features of atleast size 1 expected_graph_data["1"].x = torch.ones(2, 1) expected_graph_data["2"].x = torch.ones(1, 1) - self.assertEquals(graph_data_from_builder, expected_graph_data) + self.assertEqual(graph_data_from_builder, expected_graph_data) def test_can_create_with_preexisting_data_objects_filtering_existing_nodes_and_edges( self, @@ -158,7 +158,7 @@ def test_can_create_with_preexisting_data_objects_filtering_existing_nodes_and_e } ) - self.assertEquals(graph_data_from_builder, graph_data_1) + self.assertEqual(graph_data_from_builder, graph_data_1) # Ensure works when there are no edge features either graph_data_1["1", "1", "1"].edge_attr = None @@ -168,7 +168,7 @@ def test_can_create_with_preexisting_data_objects_filtering_existing_nodes_and_e pyg_graph_builder.add_graph_data(graph_data_1) pyg_graph_builder.add_graph_data(graph_data_2) graph_data_without_edges_from_builder = pyg_graph_builder.build() - self.assertEquals(graph_data_without_edges_from_builder, graph_data_1) + self.assertEqual(graph_data_without_edges_from_builder, graph_data_1) def test_add_subgraph_mapped_graph_data(self): pyg_graph_builder = PygGraphBuilder() @@ -260,7 +260,7 @@ def test_add_subgraph_mapped_graph_data(self): # Our expected graph does not have this since it is constructed outside the builder graph_data_from_builder.global_node_to_subgraph_node_mapping = FrozenDict({}) - self.assertEquals(graph_data_from_builder, expected_graph_data) + self.assertEqual(graph_data_from_builder, expected_graph_data) def test_feature_enforcement_policies(self): pyg_graph_builder = PygGraphBuilder() diff --git a/tests/unit/src/common/graph_builder/pyg_graph_data_test.py b/tests/unit/src/common/graph_builder/pyg_graph_data_test.py index 44d5979f6..36d17fe16 100644 --- a/tests/unit/src/common/graph_builder/pyg_graph_data_test.py +++ b/tests/unit/src/common/graph_builder/pyg_graph_data_test.py @@ -25,7 +25,7 @@ def test_equality(self): data2["1", "1", "1"].edge_index = torch.LongTensor([[0], [1]]) data2["1", "1", "2"].edge_index = torch.LongTensor([[0, 1], [0, 0]]) - self.assertEquals(data, data2) + self.assertEqual(data, data2) data = PygGraphData() data["1"].x = torch.tensor([[1, 1], [2, 2]]) @@ -39,7 +39,7 @@ def test_equality(self): data2["1"].x = torch.tensor([[1, 1], [2, 2]]) data2["2"].x = torch.tensor([[3, 3]]) - self.assertNotEquals(data, data2) + self.assertNotEqual(data, data2) data = PygGraphData() data["1"].x = torch.tensor([[1, 1], [2, 2]]) @@ -49,7 +49,7 @@ def test_equality(self): data2["1"].x = torch.tensor([[1, 1], [2, 2]]) data2["2"].x = torch.tensor([[3, 3]]) - self.assertEquals(data, data2) + self.assertEqual(data, data2) data = PygGraphData() data["1"].x = torch.tensor([[1, 1], [2, 2]]) @@ -59,4 +59,4 @@ def test_equality(self): data2["1"].x = torch.tensor([[1, 2], [2, 2]]) data2["2"].x = torch.tensor([[3, 3]]) - self.assertNotEquals(data, data2) + self.assertNotEqual(data, data2) diff --git a/tests/unit/src/training/lib/data_loaders/tf_records_iterable_dataset_test.py b/tests/unit/src/training/lib/data_loaders/tf_records_iterable_dataset_test.py index 129ef6cf3..49363744b 100644 --- a/tests/unit/src/training/lib/data_loaders/tf_records_iterable_dataset_test.py +++ b/tests/unit/src/training/lib/data_loaders/tf_records_iterable_dataset_test.py @@ -87,6 +87,6 @@ def test_loopy_iterable_dataset(self): loopy_dataset_entries = [ next(loopy_dataset_iter) for _ in range(num_records + 5) ] - self.assertEquals( + self.assertEqual( loopy_dataset_entries[0], loopy_dataset_entries[0 + num_records] ) From 1df21174598b159e701ad56ebf7e6c1b31aff1b8 Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Fri, 4 Sep 2026 00:30:57 +0000 Subject: [PATCH 2/5] Move the CUDA and Dataflow base images to Ubuntu 24.04 (glibc 2.39) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit tensorflow-data-validation 1.21.0 publishes only manylinux_2_39_x86_64 wheels, so no image below glibc 2.39 can take it. The CUDA base was Ubuntu 22.04 (glibc 2.35) and the Dataflow base was Debian bookworm (glibc 2.36). Both move here, on today's lockfile and still on Python 3.11, so that when the dependency stack moves the failure mode is attributable to the stack rather than to the OS or interpreter underneath it. - CUDA base: `nvidia/cuda:12.8.1-cudnn-devel-ubuntu24.04`, digest pinned. Drops `UV_SYSTEM_PYTHON` and the conda interpreter, so the image now builds its own `/gigl_deps/.venv` from `.python-version` exactly as the CPU base does. Verified in the built image: Python 3.11.14, torch 2.8.0+cu128, CUDA 12.8, glibc 2.39. - Dataflow base: adopts Beam's custom-container shape — own base OS, the boot harness copied from the SDK image, and an explicit ENTRYPOINT, which was previously inherited. `RUN_PYTHON_SDK_IN_DEFAULT_ENVIRONMENT=1` is required: without it boot builds a nested venv that cannot see this one, and every worker dies reporting apache-beam missing. - Dataflow images are now built for linux/amd64 only. tensorflow-data-validation ships no aarch64 wheel, so a working arm64 image was never possible; the published manifest's arm64 entry shares layer digests with amd64. - Deletes the `--inexact` branch in the installers. It existed only for the two images that set `UV_SYSTEM_PYTHON`, and neither does now. - Hardens `has_cuda_driver()` in both copies. Callers use it as an `if` condition, which suspends `set -e`, so a missing `whereis` reported "no CUDA" while the build still exited 0. That was survivable while the CUDA base shipped torch preinstalled and `--inexact` kept it; without either, it would silently install CPU torch and a WITH_CUDA=OFF GLT into a GPU image. - Adds `scripts/smoke_test_image.py`, which asserts interpreter, ABI tag, active venv, glibc floor, a caller-declared import set, and optionally CUDA, Beam version and the boot environment variable. The import set is required rather than defaulted: base images carry a metadata-only gigl-core, and the Dataflow image has no GLT by design, so a fixed list cannot describe every image. Co-Authored-By: Claude Opus 5 (1M context) --- containers/Dockerfile.cuda.base | 36 ++- containers/Dockerfile.dataflow.base | 39 ++- gigl/scripts/install_glt.sh | 11 + pyproject.toml | 11 +- requirements/install_py_deps.sh | 32 +-- scripts/build_and_push_docker_image.py | 2 +- scripts/smoke_test_image.py | 344 +++++++++++++++++++++++++ 7 files changed, 437 insertions(+), 38 deletions(-) create mode 100644 scripts/smoke_test_image.py diff --git a/containers/Dockerfile.cuda.base b/containers/Dockerfile.cuda.base index f6d05397f..7f1ab98b1 100644 --- a/containers/Dockerfile.cuda.base +++ b/containers/Dockerfile.cuda.base @@ -1,16 +1,30 @@ # syntax=docker/dockerfile:1 -FROM pytorch/pytorch:2.8.0-cuda12.8-cudnn9-devel +# The `-devel` variant ships nvcc at /usr/local/cuda/bin/nvcc, which both gigl-core's +# CMake build and GraphLearn-for-PyTorch's CUDA build need while installing dependencies; +# a runtime-only CUDA image cannot compile them. +# Pinned by digest because the version tag is floating: NVIDIA rebuilds it in place, +# and it names no cuDNN version, so the tag alone does not identify what ships here. +FROM nvidia/cuda:12.8.1-cudnn-devel-ubuntu24.04@sha256:24c8e3581ea6330038b0d374920721983312627f8adbfcf390bdb4b399d280ed + +# Keeping this image's CUDA version matched to the cu128 torch wheels pinned in uv.lock +# is load-bearing, not cosmetic. The base exports LD_LIBRARY_PATH=/usr/local/cuda/lib64, +# and glibc searches LD_LIBRARY_PATH ahead of DT_RUNPATH, which is how the nvidia-*-cu12 +# wheels point at each other. So these load out of the image rather than site-packages: +# - libcublas.so.12, libcublasLt.so.12 (via the cusolver and cublas wheels' RUNPATH) +# - libcufft.so.11 (via the cufft wheel's RUNPATH) +# - libcufile.so.0 (via the cufile wheel's RUNPATH) +# - libcusparse.so.12 (via the cusolver wheel's RUNPATH) +# - libnvJitLink.so.12 (torch's RPATH omits nvjitlink entirely) +# torch's own libraries use DT_RPATH, which outranks LD_LIBRARY_PATH, so they are +# unaffected. Move torch off cu128 without moving this base and the libraries above go +# stale, failing at `import torch` with an undefined symbol that points nowhere near +# this line. SHELL ["/bin/bash", "-c"] ENV DEBIAN_FRONTEND=noninteractive -# Already has python 3.11 installed - no need to install it again. -# We use system python since it has packages pre-installed for us. -ENV UV_SYSTEM_PYTHON=true -ENV UV_PROJECT_ENVIRONMENT=/opt/conda/ - # Install basic dependencies # TODO(mkolodner-sc): iputils-ping temporarily needed to setup inter-job VAI communication for GLT Inference. # Once VAI natively supports this communication, we can remove this requirement. @@ -31,6 +45,7 @@ COPY pyproject.toml pyproject.toml COPY uv.lock uv.lock COPY requirements requirements COPY gigl/scripts gigl/scripts +COPY .python-version .python-version # gigl-core is a path dependency in pyproject.toml. uv sync needs its metadata to # resolve the lockfile. Copying only the build manifest (no C++ sources) so cmake # configures but compiles nothing — the src Dockerfile installs the real wheel later. @@ -40,8 +55,13 @@ COPY gigl-core/README.md gigl-core/README.md RUN bash ./requirements/install_py_deps.sh -# Note: Since we are using system python, we dont need to create a virtual environment here unlike the -# cpu base image. +# The UV_PROJECT_ENVIRONMENT environment variable can be used to configure the project virtual environment path +# Since the above command should have created the .venv, we activate by default for any future uv commands. +# We also need to set VIRTUAL_ENV so pip envocations can find the virtual environment. +ENV UV_PROJECT_ENVIRONMENT=/gigl_deps/.venv +ENV VIRTUAL_ENV="${UV_PROJECT_ENVIRONMENT}" +# We just created a virtual environment, lets add the bin to the path +ENV PATH="${UV_PROJECT_ENVIRONMENT}/bin:${PATH}" # We also need to make UV detectable by the system ENV PATH="/root/.local/bin:${PATH}" diff --git a/containers/Dockerfile.dataflow.base b/containers/Dockerfile.dataflow.base index 7eb98df8d..d92bd8f50 100644 --- a/containers/Dockerfile.dataflow.base +++ b/containers/Dockerfile.dataflow.base @@ -1,10 +1,15 @@ -FROM apache/beam_python3.11_sdk:2.56.0 +# This reference is coupled twice over, and the two halves live in different places: +# - the tag (`2.56.0`) must match the `apache-beam` pin in pyproject.toml, because Beam +# requires a worker's installed apache-beam to be the same version as the submitting +# environment's, and the boot harness copied out of this stage is that entrypoint. +# - the Python minor in the image *name* (`beam_python3.11_sdk`) must match +# .python-version, which is the interpreter uv installs into /gigl_deps/.venv below. +ARG BEAM_SDK_IMAGE=apache/beam_python3.11_sdk:2.56.0 +FROM ${BEAM_SDK_IMAGE} AS beam_sdk -ENV DEBIAN_FRONTEND=noninteractive +FROM ubuntu:noble-20251001 -# We use system python for dataflow images since it has python and apache beam pre-installed. -ENV UV_SYSTEM_PYTHON=true -ENV UV_PROJECT_ENVIRONMENT=/usr/local +ENV DEBIAN_FRONTEND=noninteractive # TODO(mkolodner-sc): iputils-ping temporarily needed to setup inter-job VAI communication for GLT Inference. # Once VAI natively supports this communication, we can remove this requirement. @@ -25,6 +30,7 @@ COPY pyproject.toml pyproject.toml COPY uv.lock uv.lock COPY requirements requirements COPY gigl/scripts gigl/scripts +COPY .python-version .python-version # gigl-core is a path dependency in pyproject.toml. uv sync needs its metadata to # resolve the lockfile. Copying only the build manifest (no C++ sources) so cmake # configures but compiles nothing — the src Dockerfile installs the real wheel later. @@ -34,9 +40,28 @@ COPY gigl-core/README.md gigl-core/README.md RUN bash ./requirements/install_py_deps.sh --skip-glt-post-install -# Note: Since we are using system python, we dont need to create a virtual environment here unlike the -# cpu base image. +# The UV_PROJECT_ENVIRONMENT environment variable can be used to configure the project virtual environment path +# Since the above command should have created the .venv, we activate by default for any future uv commands. +# We also need to set VIRTUAL_ENV so pip envocations can find the virtual environment. +ENV UV_PROJECT_ENVIRONMENT=/gigl_deps/.venv +ENV VIRTUAL_ENV="${UV_PROJECT_ENVIRONMENT}" +# We just created a virtual environment, lets add the bin to the path +ENV PATH="${UV_PROJECT_ENVIRONMENT}/bin:${PATH}" # We also need to make UV detectable by the system ENV PATH="/root/.local/bin:${PATH}" +# Required, not optional. Unset, boot creates a per-worker venv with +# `python -m venv --system-site-packages` from the `python` on PATH, then requires +# apache-beam to import inside it. Seeded from /gigl_deps/.venv/bin/python, that venv +# inherits the *base* interpreter's site-packages rather than this venv's, so the import +# fails and every worker dies with "Apache Beam is not installed in the runtime +# environment". Set, boot runs the SDK directly in the environment on PATH. +ENV RUN_PYTHON_SDK_IN_DEFAULT_ENVIRONMENT=1 + WORKDIR / + +# boot is the Dataflow worker harness. It is a statically linked Go binary, so it runs on +# this rootfs unaltered even though the SDK image is built on a different distro. +# Dockerfile.dataflow.src declares no ENTRYPOINT, so this one is what workers execute. +COPY --from=beam_sdk /opt/apache/beam /opt/apache/beam +ENTRYPOINT ["/opt/apache/beam/boot"] diff --git a/gigl/scripts/install_glt.sh b/gigl/scripts/install_glt.sh index f044c6d0c..2c4aaa866 100755 --- a/gigl/scripts/install_glt.sh +++ b/gigl/scripts/install_glt.sh @@ -9,6 +9,17 @@ is_running_on_mac() { } has_cuda_driver() { + # Callers use this as an `if` condition, which suspends `set -e` for the whole + # function body. A missing `whereis` would therefore leave $cuda_location empty and + # report "no CUDA" with the build still exiting 0, installing CPU torch and a + # WITH_CUDA=OFF GLT into a CUDA image. `exit` is the only way out of a + # suspended-errexit context, so fail the build outright instead of returning. + if ! command -v whereis &> /dev/null + then + echo "whereis is unavailable, so CUDA presence cannot be determined." >&2 + exit 1 + fi + # Use the whereis command to locate the CUDA driver cuda_location=$(whereis cuda) diff --git a/pyproject.toml b/pyproject.toml index 974f409e8..249bb1321 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -47,7 +47,10 @@ dependencies = [ [project.optional-dependencies] transform = [ - # Apache Beam version has to match with what is specified in `containers/Dockerfile.dataflow.base` + # Has to match the tag in the `BEAM_SDK_IMAGE` arg default of + # `containers/Dockerfile.dataflow.base`, which supplies the Dataflow worker boot harness. + # That reference also encodes the Python minor in its image name, which tracks + # `.python-version`, so a Python upgrade moves the same line. "apache-beam[gcp]==2.56.0", "pyarrow==10.0.1", # Tensorflow transform packages. @@ -57,8 +60,10 @@ transform = [ "tensorflow-data-validation==1.16.1 ; sys_platform != 'darwin'", "tensorflow-metadata==1.16.1 ; sys_platform != 'darwin'", "tensorflow-transform==1.16.0 ; sys_platform != 'darwin'", - # Oct 28, 2025 - Latest ver tfx-bsl==1.17.1 requires latest version of Linux distros using glibc 2.39+ - # i.e. Ubuntu 24.04. Currently, not compatible with our docker images. + # Held one release behind the latest tfx-bsl==1.17.1, which needs glibc 2.39+. All three + # base images are Ubuntu 24.04 (glibc 2.39), so the OS floor is met; moving the pin is a + # dependency bump, which has to move tensorflow-data-validation, tensorflow-metadata and + # tensorflow-transform in step per the compatibility matrix linked above. "tfx-bsl==1.16.1 ; sys_platform != 'darwin'", ] pyg27-torch28-cpu = [ diff --git a/requirements/install_py_deps.sh b/requirements/install_py_deps.sh index 93a9ff0ea..ef6d61c41 100644 --- a/requirements/install_py_deps.sh +++ b/requirements/install_py_deps.sh @@ -30,6 +30,17 @@ done ### Helper functions ### has_cuda_driver() { + # Callers use this as an `if` condition, which suspends `set -e` for the whole + # function body. A missing `whereis` would therefore leave $cuda_location empty and + # report "no CUDA" with the build still exiting 0, installing CPU torch and a + # WITH_CUDA=OFF GLT into a CUDA image. `exit` is the only way out of a + # suspended-errexit context, so fail the build outright instead of returning. + if ! command -v whereis &> /dev/null + then + echo "whereis is unavailable, so CUDA presence cannot be determined." >&2 + exit 1 + fi + # Use the whereis command to locate the CUDA driver cuda_location=$(whereis cuda) @@ -121,23 +132,6 @@ install_gigl_lib_deps() { extra_deps_clause+=(--extra "$dep") done - flag_use_inexact_match="" - # If we are using system python, we want to use inexact match for dependencies so we don't override system packages. - # We currently, only do this in our Dockerfile.cuda.base, and Dockerfile.dataflow.base images, as they have python - # pre-installed with relevant packages - and currently reinstalling these in the trivial manner in a new virtual - # environment is not able to correctly symlink existing optimizations / bootstrap scripts available in the parent - # docker images of our base images. - if [[ "${UV_SYSTEM_PYTHON}" == "true" ]] - then - echo "Recognized using system python due to UV_SYSTEM_PYTHON = true." - echo "Will use inexact match for dependencies so we don't override system packages." - # Syncing is "exact" by default, which means it will remove any packages that are not present in the lockfile. - # To retain extraneous packages, use the --inexact option: - # https://docs.astral.sh/uv/concepts/projects/sync/#retaining-extraneous-packages - # This is useful for example when we might have packages pre-installed i.e. torch, pyg, etc. - flag_use_inexact_match="--inexact" - fi - # gigl-core's CMake build requires torch to be present during the build, but # torch is a runtime dep of gigl, not a declared build dep of gigl-core. # uv sync may ignore no-build-isolation-package for workspace members and always @@ -148,9 +142,9 @@ install_gigl_lib_deps() { if [[ $DEV -eq 1 ]] then # https://docs.astral.sh/uv/reference/cli/#uv-sync - uv sync ${extra_deps_clause[@]} --group dev --locked ${flag_use_inexact_match} --no-install-package gigl-core + uv sync ${extra_deps_clause[@]} --group dev --locked --no-install-package gigl-core else - uv sync ${extra_deps_clause[@]} --group gigl-core-build-backend --locked ${flag_use_inexact_match} --no-install-package gigl-core + uv sync ${extra_deps_clause[@]} --group gigl-core-build-backend --locked --no-install-package gigl-core fi uv pip install --no-build-isolation ./gigl-core/ diff --git a/scripts/build_and_push_docker_image.py b/scripts/build_and_push_docker_image.py index 4f9bd1411..85d9e7e12 100644 --- a/scripts/build_and_push_docker_image.py +++ b/scripts/build_and_push_docker_image.py @@ -92,7 +92,7 @@ def build_and_push_dataflow_image( image_name=image_name, dockerfile_path=dockerfile_path, context_path=CONTEXT_PATH, - multi_arch=True, + multi_arch=False, ) diff --git a/scripts/smoke_test_image.py b/scripts/smoke_test_image.py new file mode 100644 index 000000000..dddb7674e --- /dev/null +++ b/scripts/smoke_test_image.py @@ -0,0 +1,344 @@ +"""Asserts that a Python environment is the one a GiGL image is supposed to ship. + +Runs inside the base images, which contain no GiGL source, so the only imports at module +scope are the standard library. The packages to assert are named by the caller through +``--imports`` and imported dynamically; importing ``gigl`` here would make the script +unusable for the images it exists to check. + +The caller declares the import set because the images do not ship the same packages: + + - Base images install ``gigl-core`` from a build manifest carrying no C++ sources, so + ``gigl_core`` is a metadata-only distribution there and does not import. Only a + ``src`` image can assert it. + - The Dataflow base skips the GraphLearn-for-PyTorch post-install step, so + ``graphlearn_torch`` is absent from that image by design. + +Every check raises on mismatch, so a zero exit status is the only signal of success. + +Example: + Inside a CUDA base image:: + + $ python scripts/smoke_test_image.py --python 3.11 \ + --imports torch,graphlearn_torch \ + --venv-prefix /gigl_deps/.venv --min-glibc 2.39 + + Inside a Dataflow base image:: + + $ python scripts/smoke_test_image.py --python 3.11 \ + --imports torch,apache_beam \ + --venv-prefix /gigl_deps/.venv --beam 2.56.0 --boot-env --min-glibc 2.39 + + Inside a src image, or on a developer checkout where the venv path varies by + checkout and so cannot be asserted:: + + $ python scripts/smoke_test_image.py --python 3.11 \ + --imports torch,graphlearn_torch,gigl_core +""" + +import argparse +import importlib +import os +import platform +import sys +import sysconfig +from pathlib import Path + +SUPPORTED_PYTHON_VERSIONS: list[str] = ["3.11", "3.12", "3.13"] + +# The packages GiGL images install outside of GiGL's own source tree. Restricting +# --imports to this set turns a typo into a usage error instead of a failing import that +# reads like a broken image. +CHECKABLE_IMPORTS: list[str] = ["torch", "graphlearn_torch", "gigl_core", "apache_beam"] + +# Exported by the graphlearn_torch pybind11 extension only when it is compiled with +# WITH_CUDA=ON, so its absence identifies a CPU-only GLT wheel in a CUDA image. +GLT_CUDA_ONLY_SYMBOL = "cuda_stitch_sample_results" + + +def assert_equal(what: str, expected: object, actual: object) -> None: + """Raises unless ``actual`` equals ``expected``. + + Args: + what (str): Name of the checked property, used in the failure message. + expected (object): The value the image is supposed to have. + actual (object): The value the image actually has. + + Raises: + AssertionError: When the values differ. + """ + if actual != expected: + raise AssertionError(f"{what}: expected {expected!r}, got {actual!r}") + print(f"OK {what} == {actual!r}") + + +def parse_version(version: str) -> tuple[int, ...]: + """Splits a dotted numeric version into ints so it compares by value. + + Comparing versions as strings ranks ``"2.9"`` above ``"2.39"``, which would let a + glibc floor pass on an older libc than requested. + + Args: + version (str): Dotted version, e.g. ``"2.39"``. + + Returns: + tuple[int, ...]: The components, e.g. ``(2, 39)``. + + Raises: + ValueError: When any component is not an integer. + """ + return tuple(int(part) for part in version.split(".")) + + +def check_python(python_version: str) -> None: + """Asserts the running interpreter is the requested minor version, inside a venv. + + SOABI is the ABI tag of the running interpreter, so it is the tag every compiled + extension in the image has to carry to be importable. Requiring the delimiter after + the minor version rejects the free-threaded build, whose tag is ``cpython-313t-`` and + which no GiGL image is built against. + + Args: + python_version (str): Requested version as ``major.minor``, e.g. ``"3.13"``. + + Raises: + AssertionError: When the version, SOABI, or venv state does not match. + """ + major, minor = (int(part) for part in python_version.split(".")) + assert_equal("sys.version_info[:2]", (major, minor), sys.version_info[:2]) + + soabi = sysconfig.get_config_var("SOABI") + expected_soabi_prefix = f"cpython-{major}{minor}-" + if not isinstance(soabi, str) or not soabi.startswith(expected_soabi_prefix): + raise AssertionError( + f"SOABI: expected a value starting with {expected_soabi_prefix!r}, got {soabi!r}" + ) + print(f"OK SOABI == {soabi!r}") + + if sys.prefix == sys.base_prefix: + raise AssertionError( + "virtual environment: expected sys.prefix != sys.base_prefix, got both " + f"equal to {sys.prefix!r}" + ) + print(f"OK virtual environment active at {sys.prefix!r}") + + +def check_venv_prefix(venv_prefix: str) -> None: + """Asserts the running interpreter comes from inside ``venv_prefix``. + + Containment is by path component, not string prefix: a string prefix accepts a + sibling directory whose name merely starts the same way, so ``/gigl_deps/.venv`` + would be satisfied by an interpreter in ``/gigl_deps/.venv-broken``. + + Paths are normalized but not resolved. A venv's ``bin/python`` is a symlink to the + interpreter uv installed elsewhere, so resolving it walks out of the venv and no + prefix inside the venv could ever match. + + Args: + venv_prefix (str): Directory the interpreter is expected to live under. + + Raises: + AssertionError: When ``sys.executable`` lies outside ``venv_prefix``. + """ + executable = Path(os.path.abspath(sys.executable)) + expected_prefix = Path(os.path.abspath(venv_prefix)) + if not executable.is_relative_to(expected_prefix): + raise AssertionError( + f"sys.executable: expected a path under {str(expected_prefix)!r}, got {str(executable)!r}" + ) + print(f"OK sys.executable == {sys.executable!r}") + + +def check_imports(module_names: list[str]) -> None: + """Asserts each named package is importable. + + Args: + module_names (list[str]): Import names the image is required to provide. + + Raises: + ImportError: When any of them is missing or fails to load. + """ + for module_name in module_names: + importlib.import_module(module_name) + print(f"OK import {module_name}") + + +def check_cuda() -> None: + """Asserts CUDA is usable by torch and that GLT is built against it. + + Beyond the presence of the CUDA-only GLT symbol, this builds a ``CUDA``-mode graph + and initializes it, which copies the topology onto the device and therefore fails if + the runtime CUDA stack is broken rather than merely compiled in. + + Raises: + AssertionError: When CUDA is unavailable or GLT lacks its CUDA symbol. + """ + import torch + + assert_equal("torch.cuda.is_available()", True, torch.cuda.is_available()) + + import graphlearn_torch as glt + + if not hasattr(glt.py_graphlearn_torch, GLT_CUDA_ONLY_SYMBOL): + raise AssertionError( + f"graphlearn_torch.py_graphlearn_torch.{GLT_CUDA_ONLY_SYMBOL}: expected the " + "symbol to exist, got a CPU-only graphlearn_torch build" + ) + print(f"OK graphlearn_torch exports {GLT_CUDA_ONLY_SYMBOL}") + + edge_index = torch.tensor([[0, 1, 2], [1, 2, 0]]) + graph = glt.data.Graph(glt.data.Topology(edge_index=edge_index), mode="CUDA") + graph.lazy_init() + assert_equal("CUDA-mode graph edge count", 3, graph.edge_count) + + +def check_beam(beam_version: str) -> None: + """Asserts the installed Apache Beam matches the version the worker harness expects. + + Args: + beam_version (str): Exact expected value of ``apache_beam.__version__``. + + Raises: + AssertionError: When the installed version differs. + """ + import apache_beam + + assert_equal("apache_beam.__version__", beam_version, apache_beam.__version__) + + +def check_boot_env() -> None: + """Asserts the Beam boot harness is told to use the image's own environment. + + Raises: + AssertionError: When ``RUN_PYTHON_SDK_IN_DEFAULT_ENVIRONMENT`` is not ``"1"``, + which makes boot build a per-worker venv that cannot see the image venv's + site-packages. + """ + assert_equal( + "RUN_PYTHON_SDK_IN_DEFAULT_ENVIRONMENT", + "1", + os.environ.get("RUN_PYTHON_SDK_IN_DEFAULT_ENVIRONMENT"), + ) + + +def check_min_glibc(min_glibc: str) -> None: + """Asserts the running glibc is at least ``min_glibc``. + + The glibc floor is an installability constraint, not a preference: current + ``tensorflow-data-validation`` and ``tfx-bsl`` releases publish only + ``manylinux_2_39_x86_64`` wheels, so a base below that floor has no candidate to + resolve at all. Asserting it here names the OS as the cause; otherwise a base that + drops to an older Ubuntu surfaces as nothing but an unresolvable + ``tensorflow-data-validation`` requirement during ``uv sync``. + + Args: + min_glibc (str): Lowest acceptable glibc version, e.g. ``"2.39"``. + + Raises: + AssertionError: When glibc is older than requested, or when the running libc does + not report a version at all — a non-glibc libc has no bearing on this floor, + so it cannot satisfy it. + """ + libc, actual = platform.libc_ver() + if not actual: + raise AssertionError( + f"glibc: expected at least {min_glibc}, got no version from " + f"platform.libc_ver(), which reported libc {libc!r}" + ) + if parse_version(actual) < parse_version(min_glibc): + raise AssertionError(f"glibc: expected at least {min_glibc}, got {actual}") + print(f"OK glibc == {actual} (>= {min_glibc})") + + +def parse_imports(raw: str) -> list[str]: + """Parses the ``--imports`` value into the list of packages to require. + + Args: + raw (str): Comma-separated import names. + + Returns: + list[str]: The requested import names, in the order given. + + Raises: + argparse.ArgumentTypeError: When the value is empty or names a package outside + ``CHECKABLE_IMPORTS``. + """ + module_names = [part.strip() for part in raw.split(",") if part.strip()] + if not module_names: + raise argparse.ArgumentTypeError("expected at least one import name") + unknown = [name for name in module_names if name not in CHECKABLE_IMPORTS] + if unknown: + raise argparse.ArgumentTypeError( + f"unknown import name(s) {unknown}; choose from {CHECKABLE_IMPORTS}" + ) + return module_names + + +def main() -> None: + """Parses arguments and runs the requested checks. + + Raises: + AssertionError: When any check fails. + """ + parser = argparse.ArgumentParser( + description="Assert that the current environment matches a GiGL image contract." + ) + parser.add_argument( + "--python", + required=True, + choices=SUPPORTED_PYTHON_VERSIONS, + help="Python minor version the environment must be running.", + ) + parser.add_argument( + "--imports", + required=True, + type=parse_imports, + help="Comma-separated packages the environment must be able to import, chosen " + f"from {','.join(CHECKABLE_IMPORTS)}. Required: image package sets differ, so " + "the caller states which one it expects rather than the script guessing.", + ) + parser.add_argument( + "--venv-prefix", + default=None, + help="Directory the interpreter must live under, e.g. /gigl_deps/.venv. Omit " + "outside images, where the venv path varies by checkout.", + ) + parser.add_argument( + "--cuda", + action="store_true", + help="Also require a working CUDA runtime and a CUDA-enabled graphlearn_torch.", + ) + parser.add_argument( + "--beam", + default=None, + help="Exact apache_beam version the environment must have installed.", + ) + parser.add_argument( + "--boot-env", + action="store_true", + help="Also require RUN_PYTHON_SDK_IN_DEFAULT_ENVIRONMENT=1.", + ) + parser.add_argument( + "--min-glibc", + default=None, + help="Lowest glibc version the image may ship, e.g. 2.39.", + ) + args = parser.parse_args() + + check_python(python_version=args.python) + if args.venv_prefix is not None: + check_venv_prefix(venv_prefix=args.venv_prefix) + if args.min_glibc is not None: + check_min_glibc(min_glibc=args.min_glibc) + check_imports(module_names=args.imports) + if args.cuda: + check_cuda() + if args.beam is not None: + check_beam(beam_version=args.beam) + if args.boot_env: + check_boot_env() + + print("All checks passed.") + + +if __name__ == "__main__": + main() From 3099e5e7f80c6f7e0d51336fa058444b5685cf21 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Fri, 4 Sep 2026 16:01:46 +0000 Subject: [PATCH 3/5] [AUTOMATED] Update dep.vars, and other relevant files with new image names --- .github/cloud_builder/run_command_on_active_checkout.yaml | 2 +- gigl/dep_vars.env | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/.github/cloud_builder/run_command_on_active_checkout.yaml b/.github/cloud_builder/run_command_on_active_checkout.yaml index 1808472a0..673118a02 100644 --- a/.github/cloud_builder/run_command_on_active_checkout.yaml +++ b/.github/cloud_builder/run_command_on_active_checkout.yaml @@ -3,7 +3,7 @@ substitutions: options: logging: CLOUD_LOGGING_ONLY steps: - - name: us-central1-docker.pkg.dev/external-snap-ci-github-gigl/gigl-base-images/gigl-builder:7d3182eeb6446ce3e35910babba990c8e003879d.109.1 + - name: us-central1-docker.pkg.dev/external-snap-ci-github-gigl/gigl-base-images/gigl-builder:1df21174598b159e701ad56ebf7e6c1b31aff1b8.112.1 entrypoint: /bin/bash # Route sbt through Google's Maven Central mirror to avoid 429 rate limits from repo1.maven.org. # Intentionally set here (CI env) rather than in scala/.sbtopts or scala_spark35/.sbtopts to avoid diff --git a/gigl/dep_vars.env b/gigl/dep_vars.env index 11a852a40..dd01c39dd 100644 --- a/gigl/dep_vars.env +++ b/gigl/dep_vars.env @@ -1,7 +1,7 @@ # Note this file only supports static key value pairs so it can be loaded by make, bash, python, and sbt without any additional parsing. -DOCKER_LATEST_BASE_CUDA_IMAGE_NAME_WITH_TAG=us-central1-docker.pkg.dev/external-snap-ci-github-gigl/public-gigl/gigl-cuda-base:7d3182eeb6446ce3e35910babba990c8e003879d.109.1 -DOCKER_LATEST_BASE_CPU_IMAGE_NAME_WITH_TAG=us-central1-docker.pkg.dev/external-snap-ci-github-gigl/public-gigl/gigl-cpu-base:7d3182eeb6446ce3e35910babba990c8e003879d.109.1 -DOCKER_LATEST_BASE_DATAFLOW_IMAGE_NAME_WITH_TAG=us-central1-docker.pkg.dev/external-snap-ci-github-gigl/public-gigl/gigl-dataflow-base:7d3182eeb6446ce3e35910babba990c8e003879d.109.1 +DOCKER_LATEST_BASE_CUDA_IMAGE_NAME_WITH_TAG=us-central1-docker.pkg.dev/external-snap-ci-github-gigl/public-gigl/gigl-cuda-base:1df21174598b159e701ad56ebf7e6c1b31aff1b8.112.1 +DOCKER_LATEST_BASE_CPU_IMAGE_NAME_WITH_TAG=us-central1-docker.pkg.dev/external-snap-ci-github-gigl/public-gigl/gigl-cpu-base:1df21174598b159e701ad56ebf7e6c1b31aff1b8.112.1 +DOCKER_LATEST_BASE_DATAFLOW_IMAGE_NAME_WITH_TAG=us-central1-docker.pkg.dev/external-snap-ci-github-gigl/public-gigl/gigl-dataflow-base:1df21174598b159e701ad56ebf7e6c1b31aff1b8.112.1 DEFAULT_GIGL_RELEASE_SRC_IMAGE_CUDA=us-central1-docker.pkg.dev/external-snap-ci-github-gigl/public-gigl/src-cuda:0.3.1 DEFAULT_GIGL_RELEASE_SRC_IMAGE_CPU=us-central1-docker.pkg.dev/external-snap-ci-github-gigl/public-gigl/src-cpu:0.3.1 From 0c41d2857437453a651888d02978071c5c3ab1e2 Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Fri, 4 Sep 2026 19:23:42 +0000 Subject: [PATCH 4/5] Keep the GPU driver on LD_LIBRARY_PATH in the CUDA base image `libcuda.so.1` and `libnvidia-ml.so.1` are the host NVIDIA driver and ship in no image. Vertex AI and GKE bind-mount them into /usr/local/nvidia/lib{,64} at container start, so both directories are empty while the image builds. `nvidia/cuda` sets LD_LIBRARY_PATH=/usr/local/cuda/lib64 and drops them, which broke GPU two ways at once: - `fbgemm_gpu` names libcuda.so.1 in DT_NEEDED across 19 of its extensions, so `import torchrec` raised outright. Every trainer reaching `gigl.nn` died, because `gigl/nn/__init__.py` imports `gigl/nn/models.py` eagerly. - Everything else imported torch fine, reported no CUDA device, and trained on CPU while reporting success. Five of nine e2e pipelines passed that way on two T4s each. Move to `pytorch/pytorch:2.10.0-cuda12.8-cudnn9-devel`, the oldest torch tag on Ubuntu 24.04, which exports the driver paths as the previous 2.8.0 base did. Its own torch and Python 3.12 go unused: GiGL pins `requires-python = "==3.11.*"`, so install_py_deps.sh still builds /gigl_deps/.venv from .python-version, and the venv leads PATH. Set both variables here anyway rather than inherit them, because upstream shuffles them between releases: 2.11.0 already puts /usr/local/cuda/lib64 first. Verified in the built image: Ubuntu 24.04, glibc 2.39, nvcc 12.8, Python 3.11.14 from the venv, torch 2.8.0+cu128. With the driver directory left empty, `import torchrec` fails as it did in production; with the host driver mounted into it, the same image imports torchrec and torch 2.8.0+cu128. `smoke_test_image.py` gains `torchrec` as a checkable import, so the broken chain can be asserted directly, and `--require-nvidia-driver-path`, which compares LD_LIBRARY_PATH entries whole. That check needs no GPU, so it can gate an image build on any machine. Co-Authored-By: Claude Opus 5 (1M context) --- containers/Dockerfile.cuda.base | 41 +++++++++++---------- scripts/smoke_test_image.py | 64 ++++++++++++++++++++++++++++++++- 2 files changed, 86 insertions(+), 19 deletions(-) diff --git a/containers/Dockerfile.cuda.base b/containers/Dockerfile.cuda.base index 7f1ab98b1..b09a93c8b 100644 --- a/containers/Dockerfile.cuda.base +++ b/containers/Dockerfile.cuda.base @@ -3,28 +3,33 @@ # The `-devel` variant ships nvcc at /usr/local/cuda/bin/nvcc, which both gigl-core's # CMake build and GraphLearn-for-PyTorch's CUDA build need while installing dependencies; -# a runtime-only CUDA image cannot compile them. -# Pinned by digest because the version tag is floating: NVIDIA rebuilds it in place, -# and it names no cuDNN version, so the tag alone does not identify what ships here. -FROM nvidia/cuda:12.8.1-cudnn-devel-ubuntu24.04@sha256:24c8e3581ea6330038b0d374920721983312627f8adbfcf390bdb4b399d280ed - -# Keeping this image's CUDA version matched to the cu128 torch wheels pinned in uv.lock -# is load-bearing, not cosmetic. The base exports LD_LIBRARY_PATH=/usr/local/cuda/lib64, -# and glibc searches LD_LIBRARY_PATH ahead of DT_RUNPATH, which is how the nvidia-*-cu12 -# wheels point at each other. So these load out of the image rather than site-packages: -# - libcublas.so.12, libcublasLt.so.12 (via the cusolver and cublas wheels' RUNPATH) -# - libcufft.so.11 (via the cufft wheel's RUNPATH) -# - libcufile.so.0 (via the cufile wheel's RUNPATH) -# - libcusparse.so.12 (via the cusolver wheel's RUNPATH) -# - libnvJitLink.so.12 (torch's RPATH omits nvjitlink entirely) -# torch's own libraries use DT_RPATH, which outranks LD_LIBRARY_PATH, so they are -# unaffected. Move torch off cu128 without moving this base and the libraries above go -# stale, failing at `import torch` with an undefined symbol that points nowhere near -# this line. +# a runtime-only CUDA image cannot compile them. Its CUDA 12.8 toolkit has to stay matched +# to the cu128 torch wheels pinned in uv.lock, because those two builds compile against +# the toolkit here and link the wheels' runtime. +# The image's own torch and Python 3.12 go unused: GiGL is `requires-python = "==3.11.*"`, +# so requirements/install_py_deps.sh builds its own venv at /gigl_deps/.venv on the +# interpreter named in .python-version, and the ENV PATH below puts that venv's bin ahead +# of the image's Python. +# Pinned by digest because the version tag is floating: upstream rebuilds it in place. +# The tag has a single-platform (linux/amd64) manifest, so this digest is that manifest's, +# not an index's. +FROM pytorch/pytorch:2.10.0-cuda12.8-cudnn9-devel@sha256:b574d4ccf6d8856a5d87dcadc667aa4f95dc18d337ef3a28d02b7b01897d7081 SHELL ["/bin/bash", "-c"] ENV DEBIAN_FRONTEND=noninteractive +# The NVIDIA driver — libcuda.so.1 and libnvidia-ml.so.1 — ships in no image: Vertex AI +# and GKE bind-mount it into /usr/local/nvidia/lib{,64} at container start, so both +# directories are empty while this image builds and nothing here can verify them. +# Losing them fails two ways: fbgemm_gpu names libcuda.so.1 in DT_NEEDED across 19 of its +# extensions, so `import torchrec` raises outright, while plain torch imports fine and +# reports no CUDA device, which reads as a CPU machine rather than a broken image. +# Prepended so the base's own entries survive, and set here rather than inherited because +# upstream shuffles both variables between releases — +# pytorch/pytorch:2.11.0-cuda12.8-cudnn9-devel already puts /usr/local/cuda/lib64 first. +ENV LD_LIBRARY_PATH="/usr/local/nvidia/lib:/usr/local/nvidia/lib64:${LD_LIBRARY_PATH}" +ENV PATH="/usr/local/nvidia/bin:${PATH}" + # Install basic dependencies # TODO(mkolodner-sc): iputils-ping temporarily needed to setup inter-job VAI communication for GLT Inference. # Once VAI natively supports this communication, we can remove this requirement. diff --git a/scripts/smoke_test_image.py b/scripts/smoke_test_image.py index dddb7674e..93f482bd5 100644 --- a/scripts/smoke_test_image.py +++ b/scripts/smoke_test_image.py @@ -12,6 +12,10 @@ ``src`` image can assert it. - The Dataflow base skips the GraphLearn-for-PyTorch post-install step, so ``graphlearn_torch`` is absent from that image by design. + - ``torchrec`` pulls in ``fbgemm_gpu``, whose extensions name ``libcuda.so.1`` in + DT_NEEDED. That library belongs to the host GPU driver and is mounted into the + container at start, so asserting this import is meaningful only where the driver is + present; on a driver-less host it fails whatever the image ships. Every check raises on mismatch, so a zero exit status is the only signal of success. @@ -33,6 +37,12 @@ $ python scripts/smoke_test_image.py --python 3.11 \ --imports torch,graphlearn_torch,gigl_core + + On a GPU host, where the driver is mounted and the CUDA imports can be required:: + + $ python scripts/smoke_test_image.py --python 3.11 \ + --imports torch,graphlearn_torch,torchrec \ + --cuda --require-nvidia-driver-path """ import argparse @@ -48,12 +58,22 @@ # The packages GiGL images install outside of GiGL's own source tree. Restricting # --imports to this set turns a typo into a usage error instead of a failing import that # reads like a broken image. -CHECKABLE_IMPORTS: list[str] = ["torch", "graphlearn_torch", "gigl_core", "apache_beam"] +CHECKABLE_IMPORTS: list[str] = [ + "torch", + "graphlearn_torch", + "gigl_core", + "apache_beam", + "torchrec", +] # Exported by the graphlearn_torch pybind11 extension only when it is compiled with # WITH_CUDA=ON, so its absence identifies a CPU-only GLT wheel in a CUDA image. GLT_CUDA_ONLY_SYMBOL = "cuda_stitch_sample_results" +# Where the container runtime mounts the host GPU driver, and therefore the only place +# the loader can find libcuda.so.1 and libnvidia-ml.so.1. +NVIDIA_DRIVER_LIB_DIRS: list[str] = ["/usr/local/nvidia/lib", "/usr/local/nvidia/lib64"] + def assert_equal(what: str, expected: object, actual: object) -> None: """Raises unless ``actual`` equals ``expected``. @@ -191,6 +211,38 @@ def check_cuda() -> None: assert_equal("CUDA-mode graph edge count", 3, graph.edge_count) +def check_nvidia_driver_path() -> None: + """Asserts the GPU driver mount points are on ``LD_LIBRARY_PATH``. + + Passes on a machine with no GPU and no such directories, which is the point: they are + empty until the container runtime mounts the driver into them, so this asserts only + that the loader is pointed at them and runs as a build-time gate. Without them torch + imports fine and reports no CUDA device, so the failure otherwise surfaces as a job + that quietly trains on CPU. + + Entries are compared whole after splitting on ``os.pathsep``. A substring test would + also accept a path that merely contains one, such as + ``/opt/vendor/usr/local/nvidia/lib64x``. + + Raises: + AssertionError: When either directory is missing from ``LD_LIBRARY_PATH``. + """ + entries = [ + entry + for entry in os.environ.get("LD_LIBRARY_PATH", "").split(os.pathsep) + if entry + ] + missing = [ + directory for directory in NVIDIA_DRIVER_LIB_DIRS if directory not in entries + ] + if missing: + raise AssertionError( + f"LD_LIBRARY_PATH: expected entries {NVIDIA_DRIVER_LIB_DIRS}, got {entries} " + f"(missing {missing})" + ) + print(f"OK LD_LIBRARY_PATH contains {NVIDIA_DRIVER_LIB_DIRS}") + + def check_beam(beam_version: str) -> None: """Asserts the installed Apache Beam matches the version the worker harness expects. @@ -307,6 +359,12 @@ def main() -> None: action="store_true", help="Also require a working CUDA runtime and a CUDA-enabled graphlearn_torch.", ) + parser.add_argument( + "--require-nvidia-driver-path", + action="store_true", + help="Also require the GPU driver mount points to be on LD_LIBRARY_PATH. Needs " + "no GPU, so it is usable as a build-time gate.", + ) parser.add_argument( "--beam", default=None, @@ -329,6 +387,10 @@ def main() -> None: check_venv_prefix(venv_prefix=args.venv_prefix) if args.min_glibc is not None: check_min_glibc(min_glibc=args.min_glibc) + # Ahead of the imports: torchrec fails on a missing libcuda.so.1, and naming the + # unset path first explains why. + if args.require_nvidia_driver_path: + check_nvidia_driver_path() check_imports(module_names=args.imports) if args.cuda: check_cuda() From 42bda03b39a8ec19bc2dc83b874070f3429bbac6 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Sat, 5 Sep 2026 00:48:06 +0000 Subject: [PATCH 5/5] [AUTOMATED] Update dep.vars, and other relevant files with new image names --- .github/cloud_builder/run_command_on_active_checkout.yaml | 2 +- gigl/dep_vars.env | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/.github/cloud_builder/run_command_on_active_checkout.yaml b/.github/cloud_builder/run_command_on_active_checkout.yaml index 673118a02..93134df77 100644 --- a/.github/cloud_builder/run_command_on_active_checkout.yaml +++ b/.github/cloud_builder/run_command_on_active_checkout.yaml @@ -3,7 +3,7 @@ substitutions: options: logging: CLOUD_LOGGING_ONLY steps: - - name: us-central1-docker.pkg.dev/external-snap-ci-github-gigl/gigl-base-images/gigl-builder:1df21174598b159e701ad56ebf7e6c1b31aff1b8.112.1 + - name: us-central1-docker.pkg.dev/external-snap-ci-github-gigl/gigl-base-images/gigl-builder:0c41d2857437453a651888d02978071c5c3ab1e2.113.1 entrypoint: /bin/bash # Route sbt through Google's Maven Central mirror to avoid 429 rate limits from repo1.maven.org. # Intentionally set here (CI env) rather than in scala/.sbtopts or scala_spark35/.sbtopts to avoid diff --git a/gigl/dep_vars.env b/gigl/dep_vars.env index dd01c39dd..f7cf9df58 100644 --- a/gigl/dep_vars.env +++ b/gigl/dep_vars.env @@ -1,7 +1,7 @@ # Note this file only supports static key value pairs so it can be loaded by make, bash, python, and sbt without any additional parsing. -DOCKER_LATEST_BASE_CUDA_IMAGE_NAME_WITH_TAG=us-central1-docker.pkg.dev/external-snap-ci-github-gigl/public-gigl/gigl-cuda-base:1df21174598b159e701ad56ebf7e6c1b31aff1b8.112.1 -DOCKER_LATEST_BASE_CPU_IMAGE_NAME_WITH_TAG=us-central1-docker.pkg.dev/external-snap-ci-github-gigl/public-gigl/gigl-cpu-base:1df21174598b159e701ad56ebf7e6c1b31aff1b8.112.1 -DOCKER_LATEST_BASE_DATAFLOW_IMAGE_NAME_WITH_TAG=us-central1-docker.pkg.dev/external-snap-ci-github-gigl/public-gigl/gigl-dataflow-base:1df21174598b159e701ad56ebf7e6c1b31aff1b8.112.1 +DOCKER_LATEST_BASE_CUDA_IMAGE_NAME_WITH_TAG=us-central1-docker.pkg.dev/external-snap-ci-github-gigl/public-gigl/gigl-cuda-base:0c41d2857437453a651888d02978071c5c3ab1e2.113.1 +DOCKER_LATEST_BASE_CPU_IMAGE_NAME_WITH_TAG=us-central1-docker.pkg.dev/external-snap-ci-github-gigl/public-gigl/gigl-cpu-base:0c41d2857437453a651888d02978071c5c3ab1e2.113.1 +DOCKER_LATEST_BASE_DATAFLOW_IMAGE_NAME_WITH_TAG=us-central1-docker.pkg.dev/external-snap-ci-github-gigl/public-gigl/gigl-dataflow-base:0c41d2857437453a651888d02978071c5c3ab1e2.113.1 DEFAULT_GIGL_RELEASE_SRC_IMAGE_CUDA=us-central1-docker.pkg.dev/external-snap-ci-github-gigl/public-gigl/src-cuda:0.3.1 DEFAULT_GIGL_RELEASE_SRC_IMAGE_CPU=us-central1-docker.pkg.dev/external-snap-ci-github-gigl/public-gigl/src-cpu:0.3.1