diff --git a/dgf/src/api/BUILD b/dgf/src/api/BUILD index 7f11e67..5fe5a06 100644 --- a/dgf/src/api/BUILD +++ b/dgf/src/api/BUILD @@ -147,6 +147,7 @@ py_library( "//dgf/src/sampling:in_memory_sampler", "//dgf/src/sampling:temporal", "//dgf/src/sampling/gcp:spanner_graph_sampler", + "//dgf/src/sampling/offline_distributed:offline_distributed_gcp", ], ) diff --git a/dgf/src/api/sampling.py b/dgf/src/api/sampling.py index 1630335..bba4c2f 100644 --- a/dgf/src/api/sampling.py +++ b/dgf/src/api/sampling.py @@ -33,3 +33,4 @@ from dgf.src.sampling.gcp.spanner_graph_sampler import create_graph_spanner_sampler from dgf.src.sampling.gcp.spanner_graph_sampler import SpannerGraphSampler +from dgf.src.sampling.offline_distributed.offline_distributed_gcp import offline_distributed_sampler_gcp diff --git a/dgf/src/io/BUILD b/dgf/src/io/BUILD index 69d055a..bf21fe7 100644 --- a/dgf/src/io/BUILD +++ b/dgf/src/io/BUILD @@ -252,8 +252,6 @@ py_library( "//dgf/src/data:schema", "//dgf/src/util/weak_dep:weak_dep_base", # numpy dep, - "//third_party/py/torch:pytorch", - "//third_party/py/torch_geometric", ], ) diff --git a/dgf/src/sampling/offline_distributed/BUILD b/dgf/src/sampling/offline_distributed/BUILD new file mode 100644 index 0000000..b3c15f0 --- /dev/null +++ b/dgf/src/sampling/offline_distributed/BUILD @@ -0,0 +1,66 @@ +load("@rules_python//python:py_library.bzl", "py_library") +load("@rules_python//python:py_test.bzl", "py_test") +load("@rules_python//python:py_binary.bzl", "py_binary") + +package( + + default_visibility = [ + "//:internal_users", + "//dgf/src:__subpackages__", + ], +) + +# Libraries +# ========= + +py_library( + name = "offline_distributed_gcp", + srcs = ["offline_distributed_gcp.py"], + deps = [ + "//dgf/src/data:schema", + "//dgf/src/io:schema", + "//dgf/src/sampling:config", + "//dgf/src/util:filesystem", + "//dgf/src/util:log", + "//third_party/py/google/cloud/aiplatform", + # tqdm dep, + ], +) + +# Test +# ========= + +py_test( + name = "offline_distributed_gcp_test", + srcs = ["offline_distributed_gcp_test.py"], + deps = [ + ":offline_distributed_gcp", + # absl/testing:absltest dep, + "//dgf/src/sampling:config", + "//dgf/src/util:gen_test_graph", + "//dgf/src/util:log", + "//third_party/py/google/cloud/aiplatform", + ], +) + +# Binaries +# ======== + +py_binary( + name = "offline_distributed_gcp_integration_test", + srcs = ["offline_distributed_gcp_integration_test.py"], + tags = [ + "manual", + "notap", + ], + deps = [ + ":offline_distributed_gcp", + # absl/flags dep, + # absl/testing:absltest dep, + "//dgf/src/io:schema", + "//dgf/src/io:tf_graph_sample", + "//dgf/src/sampling:config", + "//dgf/src/util:filesystem", + "//dgf/src/util:log", + ], +) diff --git a/dgf/src/sampling/offline_distributed/__init__.py b/dgf/src/sampling/offline_distributed/__init__.py new file mode 100644 index 0000000..15e6148 --- /dev/null +++ b/dgf/src/sampling/offline_distributed/__init__.py @@ -0,0 +1,17 @@ +# Copyright 2022 Google LLC. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Offline distributed sampling on GCP.""" + +from dgf.src.sampling.offline_distributed.offline_distributed_gcp import offline_distributed_sampler_gcp diff --git a/dgf/src/sampling/offline_distributed/offline_distributed_gcp.py b/dgf/src/sampling/offline_distributed/offline_distributed_gcp.py new file mode 100644 index 0000000..af52ac5 --- /dev/null +++ b/dgf/src/sampling/offline_distributed/offline_distributed_gcp.py @@ -0,0 +1,419 @@ +# Copyright 2022 Google LLC. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Runs the distributed graph sampler on GCP using Vertex AI and Dataflow.""" + +import datetime +import os +import subprocess +import time +from typing import Any, Optional, Union + +from dgf.src.data import schema as schema_lib +from dgf.src.io import schema as io_schema +from dgf.src.sampling import config as config_lib +from dgf.src.util import filesystem +from dgf.src.util import log +from google.cloud.aiplatform import aiplatform +import tqdm + +_DEFAULT_IMAGE_URI = ( + "us-central1-docker.pkg.dev/graph-flow/glassbox-repo/sampler:latest" +) +_WORKER_BINARY = "/google3/third_party/py/dgf/src/bin/google/offline_distributed_sampling_gcp" +_WORKER_MACHINE_TYPE = "n1-highmem-4" +_DEFAULT_POLL_INTERVAL_SEC = 5.0 + +_CONSOLE_JOB_URL_TEMPLATE = "https://console.cloud.google.com/vertex-ai/locations/{region}/training/{job_id}?project={project}" +_CONSOLE_JOBS_LIST_URL_TEMPLATE = "https://console.cloud.google.com/vertex-ai/training/custom-jobs?project={project}" +_LOGGING_JOB_URL_TEMPLATE = "https://console.cloud.google.com/logs/viewer?project={project}&resource=ml_job%2Fjob_id%2F{job_id}" +_LOGGING_ROOT_URL_TEMPLATE = ( + "https://console.cloud.google.com/logs/viewer?project={project}" +) + +_TERMINAL_FAILURE_STATES = { + aiplatform.gapic.JobState.JOB_STATE_FAILED, + aiplatform.gapic.JobState.JOB_STATE_CANCELLED, + aiplatform.gapic.JobState.JOB_STATE_EXPIRED, + aiplatform.gapic.JobState.JOB_STATE_PAUSED, +} + +_STAGES = [ + "1/4: Job submitted", + "2/4: Provisioning Vertex AI worker", + "3/4: Running Dataflow distributed sampler", + "4/4: Completed", +] + + +def _format_duration(seconds: float) -> str: + """Formats a duration in seconds into human-readable format.""" + sec = int(seconds) + if sec >= 60: + return f"{sec // 60}m {sec % 60:02d}s" + return f"{sec}s" + + +def _get_default_gcp_project() -> Optional[str]: + """Gets the active GCP project from environment or gcloud config.""" + for env_var in ( + "GOOGLE_CLOUD_PROJECT", + "CLOUDSDK_CORE_PROJECT", + "GCP_PROJECT", + ): + if env_val := os.environ.get(env_var): + return env_val + try: + res = subprocess.run( + ["gcloud", "config", "get-value", "project"], + capture_output=True, + text=True, + check=False, + ) + if res.returncode == 0 and res.stdout.strip(): + return res.stdout.strip() + except Exception: # pylint: disable=broad-except + pass + return None + + +def _validate_gcs_path(path: str, param_name: str) -> str: + """Validates that a path is a Google Cloud Storage path starting with 'gs://'.""" + if not path.startswith("gs://"): + msg = ( + f"{param_name} must be a Google Cloud Storage path starting with" + f" 'gs://', got: '{path}'. You can copy data to a GCS bucket using:\n " + " gcloud storage cp -r gs:///" + ) + log.error("%s", msg) + raise ValueError(msg) + return path.rstrip("/") + + +def _validate_paths(input_path: str, output_path: str) -> tuple[str, str]: + """Validates and normalizes input and output GCS paths.""" + return ( + _validate_gcs_path(input_path, "input_path"), + _validate_gcs_path(output_path, "output_path"), + ) + + +def _get_job_urls( + project: str, region: str, resource_name: str +) -> tuple[str, str]: + """Returns the Cloud Console and Cloud Logging URLs for a CustomJob.""" + job_id = resource_name.split("/")[-1] if resource_name else "" + if job_id: + console_url = _CONSOLE_JOB_URL_TEMPLATE.format( + region=region, job_id=job_id, project=project + ) + logs_url = _LOGGING_JOB_URL_TEMPLATE.format(project=project, job_id=job_id) + else: + console_url = _CONSOLE_JOBS_LIST_URL_TEMPLATE.format(project=project) + logs_url = _LOGGING_ROOT_URL_TEMPLATE.format(project=project) + return console_url, logs_url + + +def _get_job_stage_and_name(state: Any) -> tuple[int, str]: + """Maps a JobState enum or string to a (stage_index, human_readable_name).""" + state_enum = getattr(state, "name", str(state)) + state_name = state_enum.removeprefix("JOB_STATE_").capitalize() + + if ( + state == aiplatform.gapic.JobState.JOB_STATE_SUCCEEDED + or state_enum == "JOB_STATE_SUCCEEDED" + ): + return 3, state_name + if ( + state == aiplatform.gapic.JobState.JOB_STATE_RUNNING + or "RUNNING" in state_enum + ): + return 2, state_name + return 1, state_name + + +def _write_sampling_plan( + input_path: str, + output_path: str, + plan: Union[config_lib.SimpleSamplingConfig, config_lib.SamplingPlan], + schema: Optional[schema_lib.GraphSchema], +) -> str: + """Converts sampling config to plan if needed and writes it to GCS.""" + sampling_config_path = f"{output_path}/sampling_config.json" + + if isinstance(plan, config_lib.SimpleSamplingConfig): + if schema is None: + schema_path = f"{input_path}/schema.json" + log.info("Reading graph schema from %s", schema_path) + schema = io_schema.read_schema(schema_path) + plan = config_lib.simple_sampling_config_to_sampling_plan(plan, schema) + + log.info("Writing sampling plan to %s", sampling_config_path) + with filesystem.open_write(sampling_config_path) as f: + f.write(plan.to_json(indent=2)) # pyrefly: ignore[missing-attribute] + + return sampling_config_path + + +def _create_custom_job( + input_path: str, + output_path: str, + sampling_config_path: str, + project: str, + region: str, + num_workers: int, + num_seeds: Optional[int], + staging_location: str, + temp_location: str, + display_name: str, +) -> aiplatform.CustomJob: + """Builds and instantiates the Vertex AI CustomJob.""" + args = [ + f"--input_graph={input_path}", + f"--output_samples={output_path}", + f"--sampling_config={sampling_config_path}", + f"--num_seeds={num_seeds if num_seeds is not None else 0}", + "--runner=dataflow", + f"--project={project}", + f"--region={region}", + f"--worker_machine_type={_WORKER_MACHINE_TYPE}", + f"--num_workers={num_workers}", + f"--max_num_workers={num_workers}", + "--autoscaling_algorithm=NONE", + f"--staging_location={staging_location}", + f"--temp_location={temp_location}", + "--environment_type=DOCKER", + f"--sdk_container_image={_DEFAULT_IMAGE_URI}", + f"--worker_binary={_WORKER_BINARY}", + ] + + worker_pool_specs = [{ + "machine_spec": { + "machine_type": _WORKER_MACHINE_TYPE, + }, + "replica_count": 1, + "container_spec": { + "image_uri": _DEFAULT_IMAGE_URI, + "args": args, + }, + }] + + return aiplatform.CustomJob( + display_name=display_name, + worker_pool_specs=worker_pool_specs, + project=project, + location=region, + staging_bucket=staging_location, + base_output_dir=output_path, + ) + + +def _monitor_job( + job: aiplatform.CustomJob, + project: str, + region: str, + poll_interval: float, +) -> None: + """Monitors a submitted CustomJob until completion with progress reporting.""" + resource_name = getattr(job, "resource_name", "") or "" + _, logs_url = _get_job_urls(project, region, resource_name) + + pbar = tqdm.tqdm( + total=len(_STAGES), + desc=f"Distributed Sampler [{_STAGES[0]}]", + unit="stage", + ) + pbar.update(1) + + start_time = time.time() + last_state_name = None + current_stage = 1 + + try: + while True: + state = job.state + stage_idx, state_name = _get_job_stage_and_name(state) + elapsed_str = _format_duration(time.time() - start_time) + + if state_name != last_state_name: + log.info( + "Job status changed to: %s (elapsed: %s)", state_name, elapsed_str + ) + last_state_name = state_name + + if ( + state in _TERMINAL_FAILURE_STATES + or "FAILED" in state_name.upper() + or "CANCEL" in state_name.upper() + ): + pbar.close() + job_error = getattr(job, "error", None) + error_msg = ( + f"Vertex AI CustomJob '{resource_name}' failed with status" + f" {state_name}. Error: {job_error}\nView logs: {logs_url}" + ) + log.error("%s", error_msg) + raise RuntimeError(error_msg) + + if ( + state == aiplatform.gapic.JobState.JOB_STATE_SUCCEEDED + or stage_idx == 3 + ): + if current_stage < 3: + pbar.update(3 - current_stage) + pbar.set_description(f"Distributed Sampler [{_STAGES[3]}]") + pbar.set_postfix_str(f"Success in {elapsed_str}") + pbar.close() + log.info( + "CustomJob completed successfully: %s in %s", + resource_name, + elapsed_str, + ) + break + + if stage_idx > current_stage: + pbar.update(stage_idx - current_stage) + current_stage = stage_idx + + pbar.set_description(f"Distributed Sampler [{_STAGES[stage_idx]}]") + pbar.set_postfix_str(f"Status: {state_name} | Elapsed: {elapsed_str}") + pbar.refresh() + time.sleep(poll_interval) + except Exception as e: + pbar.close() + if not isinstance(e, RuntimeError): + log.error("CustomJob monitoring error: %s", e) + raise + + +def offline_distributed_sampler_gcp( + input_path: str, + output_path: str, + plan: Union[config_lib.SimpleSamplingConfig, config_lib.SamplingPlan], + schema: Optional[schema_lib.GraphSchema] = None, + *, + blocking: bool = True, + project: Optional[str] = None, + region: str = "us-central1", + num_workers: int = 5, + num_seeds: Optional[int] = None, + temp_location: Optional[str] = None, + staging_location: Optional[str] = None, + display_name: Optional[str] = None, + poll_interval: float = _DEFAULT_POLL_INTERVAL_SEC, +) -> aiplatform.CustomJob: + """Runs the offline distributed graph sampler on GCP. + + Submits a Vertex AI CustomJob running the Glassbox container image to + execute the distributed sampling pipeline on Apache Beam / Dataflow. + + Usage example: + + ```python + job = dgf.sampling.offline_distributed_sampler_gcp( + input_path="gs://my_bucket/my_graph", + output_path="gs://my_bucket/my_samples", + plan=dgf.sampling.SimpleSamplingConfig( + seed_nodeset="paper", + num_hops=2, + hop_width=10, + ), + blocking=True, + ) + ``` + + Args: + input_path: GCS path to the input GraphFlow graph directory (must start with + 'gs://'). + output_path: GCS path to write sampled graphs and schema to (must start with + 'gs://'). + plan: Sampling configuration (`SimpleSamplingConfig` or `SamplingPlan`). + schema: Optional graph schema. If None and `plan` is `SimpleSamplingConfig`, + it is read from `{input_path}/schema.json`. + blocking: If True, waits for the job to complete while logging progress. + project: GCP project ID. If None, it is resolved from the environment. + region: GCP region to run the Vertex AI job and Dataflow workers in. + num_workers: Number of Dataflow workers. + num_seeds: Optional number of seeds to sample. If None, samples all nodes. + temp_location: Optional GCS temporary directory for Dataflow. + staging_location: Optional GCS staging directory for Dataflow and Vertex AI. + display_name: Optional display name for the Vertex AI CustomJob. + poll_interval: Polling interval in seconds when `blocking=True`. + + Returns: + The `google.cloud.aiplatform.CustomJob` instance. + """ + input_path, output_path = _validate_paths(input_path, output_path) + + if project is None: + project = _get_default_gcp_project() + if not project: + raise ValueError( + "GCP project must be specified or configured in the environment." + ) + log.info("Using GCP project: %s", project) + + if temp_location is None: + temp_location = f"{output_path}_temp" + if staging_location is None: + staging_location = f"{output_path}_staging" + if display_name is None: + timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S") + display_name = f"dgf-distributed-sampler-glassbox-{timestamp}" + + start_time = time.time() + sampling_config_path = _write_sampling_plan( + input_path=input_path, + output_path=output_path, + plan=plan, + schema=schema, + ) + + job = _create_custom_job( + input_path=input_path, + output_path=output_path, + sampling_config_path=sampling_config_path, + project=project, + region=region, + num_workers=num_workers, + num_seeds=num_seeds, + staging_location=staging_location, + temp_location=temp_location, + display_name=display_name, + ) + + log.info("Submitting Vertex AI CustomJob '%s'...", display_name) + try: + job.submit() + except Exception as e: + log.error("Failed to submit CustomJob: %s", e) + raise + + resource_name = getattr(job, "resource_name", "") or "" + console_url, logs_url = _get_job_urls(project, region, resource_name) + log.info("Vertex AI CustomJob created: %s", resource_name or display_name) + log.info(" Cloud Console: %s", console_url) + log.info(" Cloud Logging: %s", logs_url) + + if blocking: + _monitor_job( + job=job, + project=project, + region=region, + poll_interval=poll_interval, + ) + total_duration = _format_duration(time.time() - start_time) + log.info("Total duration: %s", total_duration) + + return job diff --git a/dgf/src/sampling/offline_distributed/offline_distributed_gcp_integration_test.py b/dgf/src/sampling/offline_distributed/offline_distributed_gcp_integration_test.py new file mode 100644 index 0000000..fef97bc --- /dev/null +++ b/dgf/src/sampling/offline_distributed/offline_distributed_gcp_integration_test.py @@ -0,0 +1,140 @@ +# Copyright 2022 Google LLC. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +r"""Integration test for offline_distributed_sampler_gcp on GCP. + +Usage example: + + blaze build -c opt //third_party/py/dgf/src/sampling/offline_distributed:offline_distributed_gcp_integration_test && \ + blaze-bin/third_party/py/dgf/src/sampling/offline_distributed/offline_distributed_gcp_integration_test \ + --test_dir=gs://gf-experiment-gbm-test/integration_test --alsologtostderr +""" + +import datetime +import os +from absl import flags +from absl.testing import absltest +from dgf.src.io import schema as schema_io +from dgf.src.io import tf_graph_sample as tf_graph_sample_io +from dgf.src.sampling import config as config_lib +from dgf.src.sampling.offline_distributed import offline_distributed_gcp +from dgf.src.util import filesystem +from dgf.src.util import log + +FLAGS = flags.FLAGS + +flags.DEFINE_string( + "test_dir", + "gs://gf-experiment-gbm-test/integration_test", + "Base GCS directory path for test outputs.", +) +flags.DEFINE_string( + "input_graph", + "gs://gf-experiment-gbm-test/fetch_repo/ogb_mag", + "Input GCS graph path.", +) +flags.DEFINE_string( + "project", + "graphflow-experiments-49784", + "GCP Project ID to run Vertex AI and Dataflow jobs.", +) +flags.DEFINE_string( + "region", + "us-central1", + "GCP Region.", +) +flags.DEFINE_integer( + "num_workers", + 5, + "Number of Dataflow workers.", +) +flags.DEFINE_integer( + "num_seeds", + 1000, + "Number of seeds to sample.", +) + + +class OfflineDistributedGcpIntegrationTest(absltest.TestCase): + + def test_offline_distributed_sampler_gcp(self): + timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S") + output_samples = os.path.join( + FLAGS.test_dir, f"dgf_samples_integration_test_{timestamp}" + ) + + log.info("Running integration test with output at %s", output_samples) + + plan = config_lib.SimpleSamplingConfig( + seed_nodeset="paper", + num_hops=2, + hop_width=10, + ) + + job = offline_distributed_gcp.offline_distributed_sampler_gcp( + input_path=FLAGS.input_graph, + output_path=output_samples, + plan=plan, + project=FLAGS.project, + region=FLAGS.region, + num_workers=FLAGS.num_workers, + num_seeds=FLAGS.num_seeds, + blocking=True, + ) + + self.assertIsNotNone(job.resource_name) + + # Check the structure of the output directory. + sampling_config_file = os.path.join(output_samples, "sampling_config.json") + self.assertTrue( + filesystem.exists(sampling_config_file), + f"Expected {sampling_config_file} to exist on GCS.", + ) + schema_file = os.path.join(output_samples, "schema.json") + self.assertTrue( + filesystem.exists(schema_file), + f"Expected {schema_file} to exist on GCS.", + ) + sample_shards = filesystem.glob( + os.path.join(output_samples, "samples-*.tfrecord.gz") + ) + self.assertNotEmpty( + sample_shards, + f"Expected sample shard files in {output_samples}.", + ) + + # Read the graph samples + num_read_graphs = 0 + for graph in tf_graph_sample_io.read_tfgnn_graphs( + path=os.path.join(output_samples, "samples-*.tfrecord.gz"), + schema=schema_io.read_schema(schema_file), + ): + num_read_graphs += 1 + + # The distributed sampler performs scalable Bernoulli sampling of seed + # nodes, so the number of generated graphs is approximately FLAGS.num_seeds. + self.assertAlmostEqual( + num_read_graphs, FLAGS.num_seeds, delta=int(FLAGS.num_seeds * 0.1) + ) + + log.info( + "Integration test verified %d sample shards in %s", + len(sample_shards), + output_samples, + ) + log.info("Integration test finished successfully!") + + +if __name__ == "__main__": + absltest.main() diff --git a/dgf/src/sampling/offline_distributed/offline_distributed_gcp_test.py b/dgf/src/sampling/offline_distributed/offline_distributed_gcp_test.py new file mode 100644 index 0000000..a501d21 --- /dev/null +++ b/dgf/src/sampling/offline_distributed/offline_distributed_gcp_test.py @@ -0,0 +1,375 @@ +# Copyright 2022 Google LLC. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for offline_distributed_gcp.""" + +import json +from unittest import mock +from absl.testing import absltest +from google.cloud.aiplatform import aiplatform +from dgf.src.sampling import config as config_lib +from dgf.src.sampling.offline_distributed import offline_distributed_gcp +from dgf.src.util import gen_test_graph +from dgf.src.util import log + + +class OfflineDistributedGcpTest(absltest.TestCase): + + def setUp(self): + super().setUp() + self.mock_schema = gen_test_graph.generate_schema() + self.simple_plan = config_lib.SimpleSamplingConfig( + seed_nodeset="n1", + num_hops=2, + hop_width=10, + ) + + @mock.patch.object(offline_distributed_gcp.filesystem, "open_write") + @mock.patch.object(aiplatform, "CustomJob") + def test_offline_distributed_sampler_gcp_blocking_success( + self, mock_custom_job_cls, mock_open_write + ): + mock_file = mock.MagicMock() + mock_open_write.return_value.__enter__.return_value = mock_file + mock_job = mock.MagicMock() + mock_job.resource_name = ( + "projects/test-proj/locations/us-central1/customJobs/12345" + ) + mock_job.state = aiplatform.gapic.JobState.JOB_STATE_SUCCEEDED + mock_custom_job_cls.return_value = mock_job + + with log.capture_logs(log_info=True) as logs: + job = offline_distributed_gcp.offline_distributed_sampler_gcp( + input_path="gs://my_bucket/graph", + output_path="gs://my_bucket/samples", + plan=self.simple_plan, + schema=self.mock_schema, + project="test-proj", + region="us-central1", + num_workers=3, + num_seeds=500, + blocking=True, + poll_interval=0.01, + ) + + self.assertEqual(job, mock_job) + mock_custom_job_cls.assert_called_once() + _, kwargs = mock_custom_job_cls.call_args + self.assertEqual(kwargs["project"], "test-proj") + self.assertEqual(kwargs["location"], "us-central1") + self.assertEqual(kwargs["staging_bucket"], "gs://my_bucket/samples_staging") + self.assertEqual(kwargs["base_output_dir"], "gs://my_bucket/samples") + + specs = kwargs["worker_pool_specs"] + self.assertLen(specs, 1) + self.assertEqual(specs[0]["machine_spec"]["machine_type"], "n1-highmem-4") + container_args = specs[0]["container_spec"]["args"] + self.assertIn("--input_graph=gs://my_bucket/graph", container_args) + self.assertIn("--output_samples=gs://my_bucket/samples", container_args) + self.assertIn( + "--sampling_config=gs://my_bucket/samples/sampling_config.json", + container_args, + ) + self.assertIn("--num_seeds=500", container_args) + self.assertIn("--num_workers=3", container_args) + self.assertIn("--max_num_workers=3", container_args) + self.assertIn("--runner=dataflow", container_args) + + mock_job.submit.assert_called_once() + mock_open_write.assert_called_once_with( + "gs://my_bucket/samples/sampling_config.json" + ) + written_data = "".join(call.args[0] for call in mock_file.write.call_args_list) + parsed_plan = json.loads(written_data) + self.assertEqual(parsed_plan["root"]["nodeset"], "n1") + + # Verify that console and logging URLs and duration are logged + log_messages = " ".join(l.text for l in logs) + self.assertIn( + "https://console.cloud.google.com/vertex-ai/locations/us-central1/training/12345?project=test-proj", + log_messages, + ) + self.assertIn( + "https://console.cloud.google.com/logs/viewer?project=test-proj&resource=ml_job%2Fjob_id%2F12345", + log_messages, + ) + self.assertIn("Total duration:", log_messages) + + @mock.patch.object(offline_distributed_gcp.filesystem, "open_write") + @mock.patch.object(aiplatform, "CustomJob") + def test_offline_distributed_sampler_gcp_state_transitions( + self, mock_custom_job_cls, mock_open_write + ): + mock_open_write.return_value.__enter__.return_value = mock.MagicMock() + mock_job = mock.MagicMock() + mock_job.resource_name = ( + "projects/test-proj/locations/us-central1/customJobs/12345" + ) + type(mock_job).state = mock.PropertyMock( + side_effect=[ + aiplatform.gapic.JobState.JOB_STATE_PENDING, + aiplatform.gapic.JobState.JOB_STATE_RUNNING, + aiplatform.gapic.JobState.JOB_STATE_SUCCEEDED, + ] + ) + mock_custom_job_cls.return_value = mock_job + + with log.capture_logs(log_info=True) as logs: + job = offline_distributed_gcp.offline_distributed_sampler_gcp( + input_path="gs://my_bucket/graph", + output_path="gs://my_bucket/samples", + plan=self.simple_plan, + schema=self.mock_schema, + project="test-proj", + blocking=True, + poll_interval=0.001, + ) + + self.assertEqual(job, mock_job) + mock_job.submit.assert_called_once() + log_messages = " ".join(l.text for l in logs) + self.assertIn("Pending", log_messages) + self.assertIn("Running", log_messages) + self.assertIn("CustomJob completed successfully", log_messages) + + @mock.patch.object(offline_distributed_gcp.filesystem, "open_write") + @mock.patch.object(aiplatform, "CustomJob") + def test_offline_distributed_sampler_gcp_non_blocking( + self, mock_custom_job_cls, mock_open_write + ): + mock_open_write.return_value.__enter__.return_value = mock.MagicMock() + mock_job = mock.MagicMock() + mock_job.resource_name = ( + "projects/test-proj/locations/us-central1/customJobs/12345" + ) + mock_custom_job_cls.return_value = mock_job + + job = offline_distributed_gcp.offline_distributed_sampler_gcp( + input_path="gs://my_bucket/graph", + output_path="gs://my_bucket/samples", + plan=self.simple_plan, + schema=self.mock_schema, + project="test-proj", + blocking=False, + ) + + self.assertEqual(job, mock_job) + mock_job.submit.assert_called_once() + + @mock.patch.object(offline_distributed_gcp.filesystem, "open_write") + @mock.patch.object(aiplatform, "CustomJob") + def test_offline_distributed_sampler_gcp_submit_failure( + self, mock_custom_job_cls, mock_open_write + ): + mock_open_write.return_value.__enter__.return_value = mock.MagicMock() + mock_job = mock.MagicMock() + mock_job.submit.side_effect = RuntimeError("CustomJob submission error") + mock_custom_job_cls.return_value = mock_job + + with log.capture_logs() as logs: + with self.assertRaises(RuntimeError): + offline_distributed_gcp.offline_distributed_sampler_gcp( + input_path="gs://my_bucket/graph", + output_path="gs://my_bucket/samples", + plan=self.simple_plan, + schema=self.mock_schema, + project="test-proj", + blocking=True, + ) + + self.assertTrue(any(l.severity == log.Severity.ERROR for l in logs)) + + @mock.patch.object(offline_distributed_gcp.filesystem, "open_write") + @mock.patch.object(aiplatform, "CustomJob") + def test_offline_distributed_sampler_gcp_job_failure( + self, mock_custom_job_cls, mock_open_write + ): + mock_open_write.return_value.__enter__.return_value = mock.MagicMock() + mock_job = mock.MagicMock() + mock_job.resource_name = ( + "projects/test-proj/locations/us-central1/customJobs/12345" + ) + type(mock_job).state = mock.PropertyMock( + return_value=aiplatform.gapic.JobState.JOB_STATE_FAILED + ) + mock_job.error = "Out of memory error in worker" + mock_custom_job_cls.return_value = mock_job + + with log.capture_logs() as logs: + with self.assertRaises(RuntimeError): + offline_distributed_gcp.offline_distributed_sampler_gcp( + input_path="gs://my_bucket/graph", + output_path="gs://my_bucket/samples", + plan=self.simple_plan, + schema=self.mock_schema, + project="test-proj", + blocking=True, + poll_interval=0.001, + ) + + self.assertTrue(any(l.severity == log.Severity.ERROR for l in logs)) + + @mock.patch.object(offline_distributed_gcp.io_schema, "read_schema") + @mock.patch.object(offline_distributed_gcp.filesystem, "open_write") + @mock.patch.object(aiplatform, "CustomJob") + def test_plan_conversion_and_schema_loading( + self, mock_custom_job_cls, mock_open_write, mock_read_schema + ): + mock_open_write.return_value.__enter__.return_value = mock.MagicMock() + mock_job = mock.MagicMock() + mock_job.state = aiplatform.gapic.JobState.JOB_STATE_SUCCEEDED + mock_custom_job_cls.return_value = mock_job + mock_read_schema.return_value = self.mock_schema + + offline_distributed_gcp.offline_distributed_sampler_gcp( + input_path="gs://my_bucket/graph", + output_path="gs://my_bucket/samples", + plan=self.simple_plan, + schema=None, + project="test-proj", + blocking=True, + poll_interval=0.001, + ) + + mock_read_schema.assert_called_once_with("gs://my_bucket/graph/schema.json") + + @mock.patch.object(offline_distributed_gcp, "_get_default_gcp_project") + def test_missing_project_raises_error(self, mock_get_project): + mock_get_project.return_value = None + + with self.assertRaises(ValueError): + offline_distributed_gcp.offline_distributed_sampler_gcp( + input_path="gs://my_bucket/graph", + output_path="gs://my_bucket/samples", + plan=self.simple_plan, + schema=self.mock_schema, + project=None, + ) + + @mock.patch.object(offline_distributed_gcp.filesystem, "open_write") + @mock.patch.object(aiplatform, "CustomJob") + def test_custom_parameters(self, mock_custom_job_cls, mock_open_write): + mock_open_write.return_value.__enter__.return_value = mock.MagicMock() + mock_job = mock.MagicMock() + mock_custom_job_cls.return_value = mock_job + + offline_distributed_gcp.offline_distributed_sampler_gcp( + input_path="gs://my_bucket/graph", + output_path="gs://my_bucket/samples", + plan=self.simple_plan, + schema=self.mock_schema, + project="custom-proj", + region="europe-west1", + num_workers=10, + num_seeds=None, + temp_location="gs://custom/temp", + staging_location="gs://custom/staging", + display_name="my-custom-display-name", + blocking=False, + ) + + mock_custom_job_cls.assert_called_once() + _, kwargs = mock_custom_job_cls.call_args + self.assertEqual(kwargs["display_name"], "my-custom-display-name") + self.assertEqual(kwargs["project"], "custom-proj") + self.assertEqual(kwargs["location"], "europe-west1") + self.assertEqual(kwargs["staging_bucket"], "gs://custom/staging") + + container_args = kwargs["worker_pool_specs"][0]["container_spec"]["args"] + self.assertIn("--num_seeds=0", container_args) + self.assertIn("--num_workers=10", container_args) + self.assertIn("--max_num_workers=10", container_args) + self.assertIn("--region=europe-west1", container_args) + self.assertIn("--project=custom-proj", container_args) + self.assertIn("--staging_location=gs://custom/staging", container_args) + self.assertIn("--temp_location=gs://custom/temp", container_args) + + def test_invalid_input_path_raises_error(self): + with self.assertRaisesRegex( + ValueError, "input_path must be a Google Cloud Storage path" + ): + offline_distributed_gcp.offline_distributed_sampler_gcp( + input_path="/local/path/to/graph", + output_path="gs://my_bucket/samples", + plan=self.simple_plan, + schema=self.mock_schema, + project="test-proj", + ) + + def test_invalid_output_path_raises_error(self): + with self.assertRaisesRegex( + ValueError, "output_path must be a Google Cloud Storage path" + ): + offline_distributed_gcp.offline_distributed_sampler_gcp( + input_path="gs://my_bucket/graph", + output_path="/local/path/to/samples", + plan=self.simple_plan, + schema=self.mock_schema, + project="test-proj", + ) + + def test_format_duration(self): + self.assertEqual(offline_distributed_gcp._format_duration(45), "45s") + self.assertEqual(offline_distributed_gcp._format_duration(125), "2m 05s") + self.assertEqual(offline_distributed_gcp._format_duration(3600), "60m 00s") + + def test_get_job_urls(self): + console_url, logs_url = offline_distributed_gcp._get_job_urls( + project="my-proj", + region="us-central1", + resource_name="projects/123/locations/us-central1/customJobs/456", + ) + self.assertEqual( + console_url, + "https://console.cloud.google.com/vertex-ai/locations/us-central1/training/456?project=my-proj", + ) + self.assertEqual( + logs_url, + "https://console.cloud.google.com/logs/viewer?project=my-proj&resource=ml_job%2Fjob_id%2F456", + ) + + empty_console, empty_logs = offline_distributed_gcp._get_job_urls( + project="my-proj", region="us-central1", resource_name="" + ) + self.assertEqual( + empty_console, + "https://console.cloud.google.com/vertex-ai/training/custom-jobs?project=my-proj", + ) + self.assertEqual( + empty_logs, + "https://console.cloud.google.com/logs/viewer?project=my-proj", + ) + + def test_get_job_stage_and_name(self): + stage, name = offline_distributed_gcp._get_job_stage_and_name( + aiplatform.gapic.JobState.JOB_STATE_PENDING + ) + self.assertEqual(stage, 1) + self.assertEqual(name, "Pending") + + stage, name = offline_distributed_gcp._get_job_stage_and_name( + aiplatform.gapic.JobState.JOB_STATE_RUNNING + ) + self.assertEqual(stage, 2) + self.assertEqual(name, "Running") + + stage, name = offline_distributed_gcp._get_job_stage_and_name( + aiplatform.gapic.JobState.JOB_STATE_SUCCEEDED + ) + self.assertEqual(stage, 3) + self.assertEqual(name, "Succeeded") + + +if __name__ == "__main__": + absltest.main() diff --git a/examples/BUILD b/examples/BUILD index 6d9d914..198fd1a 100644 --- a/examples/BUILD +++ b/examples/BUILD @@ -116,8 +116,19 @@ py_binary( # matplotlib dep, # numpy dep, "//third_party/py/torch:pytorch", + "//third_party/py/torch_geometric", "//third_party/py/torch_geometric:data", "//third_party/py/torch_geometric:nn", # tqdm dep, ], ) + +py_binary( + name = "create_graph_samples_offline_distributed_gcp", + srcs = ["create_graph_samples_offline_distributed_gcp.py"], + deps = [ + # absl:app dep, + # absl/flags dep, + "//dgf", + ], +) diff --git a/examples/create_graph_samples_offline_distributed_gcp.py b/examples/create_graph_samples_offline_distributed_gcp.py new file mode 100644 index 0000000..ed7896e --- /dev/null +++ b/examples/create_graph_samples_offline_distributed_gcp.py @@ -0,0 +1,128 @@ +# Copyright 2022 Google LLC. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +r"""Creates a set of graph samples using the offline distributed sampler on GCP. + +The sampling pipeline runs on Google Cloud using Vertex AI and Dataflow. + +Usage example: + +```shell +blaze run -c opt //third_party/py/dgf/examples:create_graph_samples_offline_distributed_gcp -- \ + --input_graph=gs://gf-experiment-gbm-test/fetch_repo/ogb_mag \ + --output_samples=gs://gf-experiment-gbm-test/examples/ogb_mag_samples \ + --project=graphflow-experiments-49784 \ + --seed_nodeset=paper \ + --num_hops=2 \ + --hop_width=10 \ + --num_workers=5 \ + --num_seeds=1000 \ + --alsologtostderr +``` +""" + +import datetime +import os +from typing import Sequence + +from absl import app +from absl import flags +import dgf + +_INPUT_GRAPH = flags.DEFINE_string( + "input_graph", + "gs://gf-experiment-gbm-test/fetch_repo/ogb_mag", + "Path to the input GraphFlow graph directory on GCS.", +) +_OUTPUT_SAMPLES = flags.DEFINE_string( + "output_samples", + "gs://gf-experiment-gbm-test/examples/ogb_mag_samples", + "Base GCS output directory path for graph samples.", +) +_SEED_NODESET = flags.DEFINE_string( + "seed_nodeset", + "paper", + "Seed nodeset name to sample around.", +) +_NUM_HOPS = flags.DEFINE_integer( + "num_hops", + 2, + "Number of hops in the sampling plan.", +) +_HOP_WIDTH = flags.DEFINE_integer( + "hop_width", + 10, + "Number of neighbors to sample per hop.", +) +_NUM_WORKERS = flags.DEFINE_integer( + "num_workers", + 5, + "Number of Dataflow workers.", +) +_NUM_SEEDS = flags.DEFINE_integer( + "num_seeds", + 1000, + "Number of seeds to sample (0 means all nodes).", +) +_PROJECT = flags.DEFINE_string( + "project", + None, + "GCP Project ID. If None, auto-detected from environment.", +) +_REGION = flags.DEFINE_string( + "region", + "us-central1", + "GCP Region to run Vertex AI CustomJob and Dataflow.", +) +_BLOCKING = flags.DEFINE_boolean( + "blocking", + True, + "If True, wait for the job to complete while showing progress.", +) + + +def main(argv: Sequence[str]) -> None: + if len(argv) > 1: + raise app.UsageError("Too many command-line arguments.") + + timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S") + output_path = os.path.join( + _OUTPUT_SAMPLES.value, f"samples_{timestamp}" + ) + + print(f"Starting offline distributed sampling from {_INPUT_GRAPH.value} to {output_path}...") + + plan = dgf.sampling.SimpleSamplingConfig( + seed_nodeset=_SEED_NODESET.value, + num_hops=_NUM_HOPS.value, + hop_width=_HOP_WIDTH.value, + ) + + job = dgf.sampling.offline_distributed_sampler_gcp( + input_path=_INPUT_GRAPH.value, + output_path=output_path, + plan=plan, + project=_PROJECT.value, + region=_REGION.value, + num_workers=_NUM_WORKERS.value, + num_seeds=_NUM_SEEDS.value, + blocking=_BLOCKING.value, + ) + + print(f"Sampling job finished. Vertex AI Job Resource: {job.resource_name}") + print(f"Generated samples available at: {output_path}") + + +if __name__ == "__main__": + app.run(main)