Repository: awslabs/graphstorm

Issue:
Add support for EMR Serverless execution
For larger graphs (30B+ edges) the 20 instance hard limit of SageMaker Processing can be a blocker. We want to provide another execution environment that scales to even larger graphs, while avoiding any infra management (cluster spin-up/shutdown), which EMR Serverless provides.


Discussion:

Resolving patch:
diff --git a/graphstorm-processing/graphstorm_processing/data_transformations/dist_transformations/dist_label_transformation.py b/graphstorm-processing/graphstorm_processing/data_transformations/dist_transformations/dist_label_transformation.py
--- a/graphstorm-processing/graphstorm_processing/data_transformations/dist_transformations/dist_label_transformation.py
+++ b/graphstorm-processing/graphstorm_processing/data_transformations/dist_transformations/dist_label_transformation.py
@@ -59,7 +59,7 @@ def apply(self, input_df: DataFrame) -> DataFrame:
             self.label_column,
         )
 
-        # Labels that were and were assigned the value numLabels by the StringIndexer
+        # Labels that were missing and were assigned the value numLabels by the StringIndexer
         # are converted to None
         long_class_label = indexed_df.select(F.col(self.label_column).cast("long")).select(
             F.when(
diff --git a/graphstorm-processing/graphstorm_processing/data_transformations/s3_utils.py b/graphstorm-processing/graphstorm_processing/data_transformations/s3_utils.py
--- a/graphstorm-processing/graphstorm_processing/data_transformations/s3_utils.py
+++ b/graphstorm-processing/graphstorm_processing/data_transformations/s3_utils.py
@@ -43,7 +43,7 @@ def list_s3_objects(bucket: str, prefix: str, s3_boto_client: boto3.client = Non
     """
     assert not prefix.startswith(
         "s3://"
-    ), f"Prefix should not start with 's3://' but be relative to bucket, got {prefix}"
+    ), f"Prefix should not start with 's3://' but be relative to the bucket, {bucket}, got {prefix}"
     s3_boto_client = get_high_retry_s3_client() if s3_boto_client is None else s3_boto_client
     paginator = s3_boto_client.get_paginator("list_objects_v2")
     pages = paginator.paginate(Bucket=bucket, Prefix=prefix)
@@ -82,14 +82,15 @@ def extract_bucket_and_key(
     path_with_bucket: str, relative_path: Optional[str] = None
 ) -> tuple[str, str]:
     """Given an S3 path that includes a bucket, and a relative path,
-    extracts the bucket name and full key path.
+    extracts the bucket name and full key path. If only `path_with_bucket`
+    is provided, will split that path into bucket name and prefix path.
 
     Parameters
     ----------
     path_with_bucket : str
         An S3 path that can include a bucket name and a key prefix, e.g.
         's3://my-bucket/my/prefix/'.
-    relative_path : Optional[str], optional
+    relative_path : Optional[str]
         An S3 key path that's relative to `path_with_bucket`, e.g.
         'rest/of/path/to/key'. If not provided only `path_with_bucket`
         will be split.
@@ -99,9 +100,21 @@ def extract_bucket_and_key(
     str
         A tuple whose first element is the bucket name and the second
         the full path to the key.
+
+    Example
+    -------
+    .. code::
+
+        >>> extract_bucket_and_key("s3://my-bucket/prefix", "rest/of/path/to/key")
+        ("my_bucket", "prefix/rest/of/path/to/key")
+        >>> extract_bucket_and_key("s3://my-bucket/prefix/key")
+        ("my_bucket", "prefix/key")
+        >>> extract_bucket_and_key("s3://my-bucket/")
+        ("my_bucket", "")
     """
     if not path_with_bucket.startswith("s3://"):
         path_with_bucket = f"s3://{path_with_bucket}"
+    path_with_bucket = s3_path_remove_trailing(path_with_bucket)
     if relative_path:
         if relative_path.startswith("/"):
             relative_path = relative_path[1:]
@@ -111,7 +124,11 @@ def extract_bucket_and_key(
     # We split on '/' to get the bucket, as it's always the third split element in an S3 URI
     file_bucket = file_s3_uri.split("/")[2]
     # Similarly, by having maxsplit=3 we get the S3 key value as the fourth element
-    file_key = file_s3_uri.split("/", 3)[3]
+    file_parts = file_s3_uri.split("/", maxsplit=3)
+    if len(file_parts) == 4:
+        file_key = file_parts[3]
+    else:
+        file_key = ""
     # We remove any trailing '/' from the key
     file_key = s3_path_remove_trailing(file_key)
 
diff --git a/graphstorm-processing/graphstorm_processing/graph_loaders/dist_heterogeneous_loader.py b/graphstorm-processing/graphstorm_processing/graph_loaders/dist_heterogeneous_loader.py
--- a/graphstorm-processing/graphstorm_processing/graph_loaders/dist_heterogeneous_loader.py
+++ b/graphstorm-processing/graphstorm_processing/graph_loaders/dist_heterogeneous_loader.py
@@ -28,6 +28,7 @@
     StructField,
     StringType,
     IntegerType,
+    LongType,
     ArrayType,
     ByteType,
 )
@@ -122,6 +123,7 @@ def __init__(
             if num_output_files and num_output_files > 0
             else int(spark.sparkContext.defaultParallelism)
         )
+        assert self.num_output_files > 0
         # Mapping from node type to filepath, each file is a node-str to node-int-id mapping
         self.node_mapping_paths = {}  # type: Dict[str, Sequence[str]]
         # Mapping from label name to value counts
@@ -703,7 +705,7 @@ def _extend_mapping_from_edges(
             map_schema = StructType(
                 [
                     StructField(join_col, StringType(), True),
-                    StructField(index_col, IntegerType(), True),
+                    StructField(index_col, LongType(), True),
                 ]
             )
 
@@ -742,7 +744,7 @@ def create_node_id_map_from_nodes_df(self, node_df: DataFrame, node_col: str) ->
         node_rdd_with_ids = node_df.rdd.zipWithIndex()
 
         node_id_col = f"{node_col}-int_id"
-        new_schema = original_schema.add(StructField(node_id_col, IntegerType(), False))
+        new_schema = original_schema.add(StructField(node_id_col, LongType(), False))
 
         node_rdd_with_ids = node_rdd_with_ids.map(
             lambda rdd_row: (list(rdd_row[0]) + [rdd_row[1]])  # type: ignore
@@ -1064,6 +1066,11 @@ def write_edge_structure(
             The first list contains the original edge files, the second is the reversed
             edge files, will be empty if `self.add_reverse_edges` is False.
         """
+        # TODO: An option for dealing with skewed data:
+        # Find the heavy hitter, collect, do broadcast join with just them,
+        # (for both sides of the edge and filter the rows?) then do the join with
+        # the rest, then concat the two (or write to storage separately, concat at read-time)
+
         src_col = edge_config.src_col
         src_ntype = edge_config.src_ntype
         dst_col = edge_config.dst_col
@@ -1084,8 +1091,10 @@ def write_edge_structure(
             )
             .withColumnRenamed(NODE_MAPPING_INT, "src_int_id")
             .withColumnRenamed(NODE_MAPPING_STR, "src_str_id")
+            .repartition(self.num_output_files, F.col("src_str_id"))
         )
         # Join incoming edge df with mapping df to transform source str-ids to int ids
+        edge_df = edge_df.repartition(self.num_output_files, F.col(src_col))
         edge_df_with_int_src = src_node_id_mapping.join(
             edge_df,
             src_node_id_mapping["src_str_id"] == edge_df[src_col],
@@ -1115,9 +1124,13 @@ def write_edge_structure(
             )
             .withColumnRenamed(NODE_MAPPING_INT, "dst_int_id")
             .withColumnRenamed(NODE_MAPPING_STR, "dst_str_id")
+            .repartition(self.num_output_files, F.col("dst_str_id"))
         )
         # Join the newly created src-int-id edge df with mapping
         # df to transform destination str-ids to int ids
+        edge_df_with_int_src = edge_df_with_int_src.repartition(
+            self.num_output_files, F.col(dst_col)
+        )
         edge_df_with_int_ids = dst_node_id_mapping.join(
             edge_df_with_int_src,
             dst_node_id_mapping["dst_str_id"] == edge_df_with_int_src[dst_col],


Test patch:
diff --git a/graphstorm-processing/tests/test_s3_utils.py b/graphstorm-processing/tests/test_s3_utils.py
new file mode 100644
--- /dev/null
+++ b/graphstorm-processing/tests/test_s3_utils.py
@@ -0,0 +1,36 @@
+"""
+Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+
+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
+
+    http://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.
+"""
+
+from graphstorm_processing.data_transformations import s3_utils
+
+
+def test_extract_bucket_and_key():
+    """Test various conditions for extracting bucket and key"""
+    bucket, key = s3_utils.extract_bucket_and_key("s3://my-bucket/my-key")
+    assert bucket == "my-bucket"
+    assert key == "my-key"
+
+    bucket, key = s3_utils.extract_bucket_and_key("s3://my-bucket")
+    assert bucket == "my-bucket"
+    assert key == ""
+
+    bucket, key = s3_utils.extract_bucket_and_key("s3://my-bucket", "/my-key/subkey")
+    assert bucket == "my-bucket"
+    assert key == "my-key/subkey"
+
+    bucket, key = s3_utils.extract_bucket_and_key("s3://my-bucket/my-key/", "subkey/")
+    assert bucket == "my-bucket"
+    assert key == "my-key/subkey"