GCP Distributed sampler Python API

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",
     ],
 )
 

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

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",
     ],
 )
 

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",
+    ],
+)

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