[CELEBORN-2254] Fix support for S3 and add a simple integration test

### What changes were proposed in this pull request?

* Fix creating files to S3 (and other DFS)
* Add integration test for Spark and S3 (using Minio)
* in CI some job will run with the AWS profile because this way we can activate the new integration test (that needs the S3 client dependencies)

### Why are the changes needed?

See https://issues.apache.org/jira/browse/CELEBORN-2254.

### Does this PR resolve a correctness bug?

No

### Does this PR introduce _any_ user-facing change?

No

### How was this patch tested?

* I have added an integration test
* I have this patch on out internal fork, to make Celeborn run on k8s with S3

Closes apache#3592 from eolivelli/CELEBORN-2254-apache.

Authored-by: Enrico Olivelli <enrico@beast.io>
Signed-off-by: SteNicholas <programgeek@163.com>

Author: Enrico Olivelli

.github/workflows/maven.yml

@@ -133,7 +133,8 @@ jobs:
       run: |
         SPARK_BINARY_VERSION=${{ matrix.spark }}
         SPARK_MAJOR_VERSION=${SPARK_BINARY_VERSION%%.*}
-        PROFILES="-Pgoogle-mirror,spark-${{ matrix.spark }}"
+        # enable AWS profile to run end-to-end tests with S3 storage
+        PROFILES="-Pgoogle-mirror,aws,spark-${{ matrix.spark }}"
         TEST_MODULES="client-spark/common,client-spark/spark-${SPARK_MAJOR_VERSION},client-spark/spark-${SPARK_MAJOR_VERSION}-columnar-common,client-spark/spark-${SPARK_MAJOR_VERSION}-shaded,tests/spark-it"
         build/mvn $PROFILES -pl $TEST_MODULES -am clean install -DskipTests
         build/mvn $PROFILES -pl $TEST_MODULES -Dspark.shuffle.sort.io.plugin.class=${{ matrix.shuffle-plugin-class }} test
@@ -170,7 +171,8 @@ jobs:
         run: |
           SPARK_BINARY_VERSION=${{ matrix.spark }}
           SPARK_MAJOR_VERSION=${SPARK_BINARY_VERSION%%.*}
-          PROFILES="-Pgoogle-mirror,spark-${{ matrix.spark }}"
+          # enable AWS profile to run end-to-end tests with S3 storage
+          PROFILES="-Pgoogle-mirror,aws,spark-${{ matrix.spark }}"
           TEST_MODULES="client-spark/common,client-spark/spark-3,client-spark/spark-3-columnar-common,client-spark/spark-${SPARK_MAJOR_VERSION}-shaded,tests/spark-it"
           build/mvn $PROFILES -pl $TEST_MODULES -am clean install -DskipTests
           build/mvn $PROFILES -pl $TEST_MODULES -Dspark.shuffle.sort.io.plugin.class=${{ matrix.shuffle-plugin-class }} test

client/src/main/java/org/apache/celeborn/client/ShuffleClient.java

@@ -19,6 +19,7 @@
 
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.ConcurrentHashMap;
@@ -122,6 +123,7 @@ public static Map<StorageInfo.Type, FileSystem> getHadoopFs(CelebornConf conf) {
             hadoopFs = CelebornHadoopUtils.getHadoopFS(conf);
           } catch (Exception e) {
             logger.error("Celeborn initialize DFS failed.", e);
+            hadoopFs = Collections.emptyMap();
           }
         }
       }

client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java

@@ -36,6 +36,7 @@
 import io.netty.buffer.Unpooled;
 import org.apache.commons.lang3.StringUtils;
 import org.apache.commons.lang3.tuple.Pair;
+import org.apache.hadoop.fs.FileSystem;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -231,6 +232,13 @@ public ShuffleClientImpl(String appUniqueId, CelebornConf conf, UserIdentifier u
 
     reviveManager = new ReviveManager(this, conf);
 
+    if (conf.hasS3Storage()) {
+      Map<StorageInfo.Type, FileSystem> hadoopFs = getHadoopFs(conf);
+      FileSystem s3client = hadoopFs.get(StorageInfo.Type.S3);
+      logger.info("S3 client: {}", s3client);
+      if (s3client == null)
+        throw new IllegalStateException("S3 type is requred but the S3 client was not created");
+    }
     logger.info("Created ShuffleClientImpl, appUniqueId: {}", appUniqueId);
   }
 

common/src/main/scala/org/apache/celeborn/common/util/CelebornHadoopUtils.scala

@@ -99,6 +99,7 @@ object CelebornHadoopUtils extends Logging {
       dirs.foreach {
         case (storageType, dir) => {
           val path = new Path(dir)
+          logInfo(s"Creating HadoopFS for type $storageType at path $path");
           hadoopFs.put(storageType, path.getFileSystem(hadoopConf))
         }
       })

multipart-uploader/multipart-uploader-s3/src/main/java/org/apache/celeborn/S3MultipartUploadHandler.java

@@ -27,12 +27,14 @@
 import com.amazonaws.AmazonClientException;
 import com.amazonaws.ClientConfiguration;
 import com.amazonaws.auth.EnvironmentVariableCredentialsProvider;
+import com.amazonaws.client.builder.AwsClientBuilder;
 import com.amazonaws.event.ProgressListener;
 import com.amazonaws.retry.PredefinedBackoffStrategies;
 import com.amazonaws.retry.PredefinedRetryPolicies;
 import com.amazonaws.retry.RetryPolicy;
 import com.amazonaws.services.s3.AmazonS3;
 import com.amazonaws.services.s3.AmazonS3ClientBuilder;
+import com.amazonaws.services.s3.model.*;
 import com.amazonaws.services.s3.model.AbortMultipartUploadRequest;
 import com.amazonaws.services.s3.model.CompleteMultipartUploadRequest;
 import com.amazonaws.services.s3.model.CompleteMultipartUploadResult;
@@ -61,15 +63,15 @@ public class S3MultipartUploadHandler implements MultipartUploadHandler {
 
   private String uploadId;
 
-  private AmazonS3 s3Client;
+  private final AmazonS3 s3Client;
 
-  private String key;
+  private final String key;
 
-  private String bucketName;
+  private final String bucketName;
 
-  private Integer s3MultiplePartUploadMaxRetries;
-  private Integer baseDelay;
-  private Integer maxBackoff;
+  private final Integer s3MultiplePartUploadMaxRetries;
+  private final Integer baseDelay;
+  private final Integer maxBackoff;
 
   public S3MultipartUploadHandler(
       FileSystem hadoopFs,
@@ -103,12 +105,23 @@ public S3MultipartUploadHandler(
         new ClientConfiguration()
             .withRetryPolicy(retryPolicy)
             .withMaxErrorRetry(s3MultiplePartUploadMaxRetries);
-    this.s3Client =
+    AmazonS3ClientBuilder builder =
         AmazonS3ClientBuilder.standard()
             .withCredentials(providers)
-            .withRegion(conf.get(Constants.AWS_REGION))
-            .withClientConfiguration(clientConfig)
-            .build();
+            .withClientConfiguration(clientConfig);
+    // for MinIO
+    String endpoint = conf.get("fs.s3a.endpoint");
+    if (endpoint != null && !endpoint.isEmpty()) {
+      builder =
+          builder
+              .withEndpointConfiguration(
+                  new AwsClientBuilder.EndpointConfiguration(
+                      endpoint, conf.get(Constants.AWS_REGION)))
+              .withPathStyleAccessEnabled(conf.getBoolean("fs.s3a.path.style.access", false));
+    } else {
+      builder = builder.withRegion(conf.get(Constants.AWS_REGION));
+    }
+    this.s3Client = builder.build();
     this.key = key;
   }
 

pom.xml

@@ -155,6 +155,7 @@
     <maven.plugin.silencer.version>1.7.19</maven.plugin.silencer.version>
     <maven.plugin.resources.version>3.3.1</maven.plugin.resources.version>
     <openapi.generator.version>7.8.0</openapi.generator.version>
+    <testcontainers-minio.version>1.21.4</testcontainers-minio.version>
 
     <!-- Allow modules to enable / disable certain build plugins easily. -->
     <testJarPhase>prepare-package</testJarPhase>
@@ -525,6 +526,12 @@
         <version>${jakarta.ws.rs-api.version}</version>
       </dependency>
 
+      <dependency>
+        <groupId>org.testcontainers</groupId>
+        <artifactId>minio</artifactId>
+        <version>${testcontainers-minio.version}</version>
+      </dependency>
+
       <dependency>
         <groupId>org.openapitools</groupId>
         <artifactId>jackson-databind-nullable</artifactId>

project/CelebornBuild.scala

@@ -90,6 +90,9 @@ object Dependencies {
   val picocliVersion = "4.7.6"
   val jmhVersion = "1.37"
 
+  // S3 integration tests with Minio
+  val testContainerMinioVersion = "1.21.4"
+
   // For SSL support
   val bouncycastleVersion = "1.77"
 
@@ -213,6 +216,8 @@ object Dependencies {
   val jakartaAnnotationApi = "jakarta.annotation" % "jakarta.annotation-api" % jakartaAnnotationApiVersion
   val jakartaWsRsApi = "jakarta.ws.rs" % "jakarta.ws.rs-api" % jakartaWsRsApiVersion
 
+  val testContainerMinio = "org.testcontainers" % "minio" %  testContainerMinioVersion
+
   // Test dependencies
   // https://www.scala-sbt.org/1.x/docs/Testing.html
   val junitInterface = "com.github.sbt" % "junit-interface" % junitInterfaceVersion
@@ -425,7 +430,8 @@ object CelebornCommonSettings {
     Dependencies.scalatest % "test",
     Dependencies.junit % "test",
     // https://www.scala-sbt.org/1.x/docs/Testing.html
-    Dependencies.junitInterface % "test")
+    Dependencies.junitInterface % "test",
+    Dependencies.testContainerMinio  % "test")
 }
 
 object CelebornBuild extends sbt.internal.BuildDef {

tests/spark-it/pom.xml

@@ -182,6 +182,11 @@
       <artifactId>jakarta.servlet-api</artifactId>
       <scope>test</scope>
     </dependency>
+    <dependency>
+      <groupId>org.testcontainers</groupId>
+      <artifactId>minio</artifactId>
+      <scope>test</scope>
+    </dependency>
   </dependencies>
   <profiles>
     <profile>

tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/s3/BasicEndToEndTieredStorageTest.scala

@@ -0,0 +1,133 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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.
+ */
+
+package org.apache.celeborn.tests.spark
+
+import org.apache.spark.SparkConf
+import org.apache.spark.sql.SparkSession
+import org.scalatest.BeforeAndAfterEach
+import org.scalatest.funsuite.AnyFunSuite
+import org.testcontainers.containers.MinIOContainer
+
+import org.apache.celeborn.client.ShuffleClient
+import org.apache.celeborn.common.CelebornConf
+import org.apache.celeborn.common.protocol.ShuffleMode
+
+class BasicEndToEndTieredStorageTest extends AnyFunSuite
+  with SparkTestBase
+  with BeforeAndAfterEach {
+
+  var container: MinIOContainer = null;
+  val skipAWSTest = !isClassPresent("org.apache.hadoop.fs.s3a.S3AFileSystem")
+
+  def isClassPresent(className: String): Boolean = {
+    try {
+      Class.forName(className)
+      true
+    } catch {
+      case _: ClassNotFoundException => false
+    }
+  }
+
+  override def beforeAll(): Unit = {
+
+    if (skipAWSTest)
+      return
+
+    container = new MinIOContainer("minio/minio:RELEASE.2023-09-04T19-57-37Z");
+    container.start()
+
+    // create bucket using Minio command line tool
+    container.execInContainer(
+      "mc",
+      "alias",
+      "set",
+      "dockerminio",
+      "http://minio:9000",
+      container.getUserName,
+      container.getPassword)
+    container.execInContainer("mc", "mb", "dockerminio/sample-bucket")
+
+    System.setProperty("aws.accessKeyId", container.getUserName)
+    System.setProperty("aws.secretKey", container.getPassword)
+
+    val s3url = container.getS3URL
+    val augmentedConfiguration = Map(
+      CelebornConf.ACTIVE_STORAGE_TYPES.key -> "MEMORY,S3",
+      CelebornConf.WORKER_STORAGE_CREATE_FILE_POLICY.key -> "MEMORY,S3",
+      CelebornConf.WORKER_STORAGE_EVICT_POLICY.key -> "MEMORY|S3",
+      "celeborn.hadoop.fs.s3a.endpoint" -> s"$s3url",
+      "celeborn.hadoop.fs.s3a.aws.credentials.provider" -> "org.apache.hadoop.fs.s3a.SimpleAWSCredentialsProvider",
+      "celeborn.hadoop.fs.s3a.access.key" -> container.getUserName,
+      "celeborn.hadoop.fs.s3a.secret.key" -> container.getPassword,
+      "celeborn.hadoop.fs.s3a.path.style.access" -> "true",
+      CelebornConf.S3_DIR.key -> "s3://sample-bucket/test/celeborn",
+      CelebornConf.S3_ENDPOINT_REGION.key -> "dummy-region")
+
+    setupMiniClusterWithRandomPorts(
+      masterConf = augmentedConfiguration,
+      workerConf = augmentedConfiguration,
+      workerNum = 1)
+  }
+
+  override def beforeEach(): Unit = {
+    ShuffleClient.reset()
+  }
+
+  override def afterAll(): Unit = {
+    System.clearProperty("aws.accessKeyId")
+    System.clearProperty("aws.secretKey")
+    if (container != null) {
+      container.close()
+      super.afterAll()
+    }
+  }
+
+  override def updateSparkConf(sparkConf: SparkConf, mode: ShuffleMode): SparkConf = {
+    val s3url = container.getS3URL
+    val newConf = sparkConf
+      .set("spark." + CelebornConf.ACTIVE_STORAGE_TYPES.key, "MEMORY,S3")
+      .set("spark." + CelebornConf.S3_DIR.key, "s3://sample-bucket/test/celeborn")
+      .set("spark." + CelebornConf.S3_ENDPOINT_REGION.key, "dummy-region")
+      .set("spark.celeborn.hadoop.fs.s3a.endpoint", s"$s3url")
+      .set(
+        "spark.celeborn.hadoop.fs.s3a.aws.credentials.provider",
+        "org.apache.hadoop.fs.s3a.SimpleAWSCredentialsProvider")
+      .set("spark.celeborn.hadoop.fs.s3a.access.key", container.getUserName)
+      .set("spark.celeborn.hadoop.fs.s3a.secret.key", container.getPassword)
+      .set("spark.celeborn.hadoop.fs.s3a.path.style.access", "true")
+
+    super.updateSparkConf(newConf, mode)
+  }
+
+  test("celeborn spark integration test - s3") {
+    assume(
+      !skipAWSTest,
+      "Skipping test because AWS Hadoop client is not in the classpath (enable with -Paws")
+
+    val s3url = container.getS3URL
+    log.info(s"s3url $s3url");
+    val sparkConf = new SparkConf().setAppName("celeborn-demo").setMaster("local[2]")
+    val celebornSparkSession = SparkSession.builder()
+      .config(updateSparkConf(sparkConf, ShuffleMode.HASH))
+      .getOrCreate()
+    groupBy(celebornSparkSession)
+
+    celebornSparkSession.stop()
+  }
+
+}

worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/StorageManager.scala

@@ -43,6 +43,7 @@ import org.apache.celeborn.common.meta.{DeviceInfo, DiskFileInfo, DiskInfo, Disk
 import org.apache.celeborn.common.metrics.source.{AbstractSource, ThreadPoolSource}
 import org.apache.celeborn.common.network.util.{NettyUtils, TransportConf}
 import org.apache.celeborn.common.protocol.{PartitionLocation, PartitionSplitMode, PartitionType, StorageInfo}
+import org.apache.celeborn.common.protocol.StorageInfo.Type
 import org.apache.celeborn.common.quota.ResourceConsumption
 import org.apache.celeborn.common.util.{CelebornExitKind, CelebornHadoopUtils, CollectionUtils, DiskUtils, JavaUtils, PbSerDeUtils, ThreadUtils, Utils}
 import org.apache.celeborn.service.deploy.worker._
@@ -1081,6 +1082,7 @@ final private[worker] class StorageManager(conf: CelebornConf, workerSource: Abs
       partitionType: PartitionType,
       partitionSplitEnabled: Boolean): (Flusher, DiskFileInfo, File) = {
     val suggestedMountPoint = location.getStorageInfo.getMountPoint
+    val storageType = location.getStorageInfo.getType
     var retryCount = 0
     var exception: IOException = null
     val shuffleKey = Utils.makeShuffleKey(appId, shuffleId)
@@ -1102,7 +1104,7 @@ final private[worker] class StorageManager(conf: CelebornConf, workerSource: Abs
         throw new IOException(s"No available disks! suggested mountPoint $suggestedMountPoint")
       }
 
-      if (dirs.isEmpty && location.getStorageInfo.HDFSAvailable()) {
+      if (storageType == Type.HDFS && location.getStorageInfo.HDFSAvailable()) {
         val shuffleDir =
           new Path(new Path(hdfsDir, conf.workerWorkingDir), s"$appId/$shuffleId")
         FileSystem.mkdirs(
@@ -1120,9 +1122,10 @@ final private[worker] class StorageManager(conf: CelebornConf, workerSource: Abs
           fileName,
           hdfsFileInfo)
         return (hdfsFlusher.get, hdfsFileInfo, null)
-      } else if (dirs.isEmpty && location.getStorageInfo.S3Available()) {
+      } else if (storageType == Type.S3 && location.getStorageInfo.S3Available()) {
         val shuffleDir =
           new Path(new Path(s3Dir, conf.workerWorkingDir), s"$appId/$shuffleId")
+        logDebug(s"trying to create S3 file at $shuffleDir");
         FileSystem.mkdirs(
           StorageManager.hadoopFs.get(StorageInfo.Type.S3),
           shuffleDir,
@@ -1138,7 +1141,7 @@ final private[worker] class StorageManager(conf: CelebornConf, workerSource: Abs
           fileName,
           s3FileInfo)
         return (s3Flusher.get, s3FileInfo, null)
-      } else if (dirs.isEmpty && location.getStorageInfo.OSSAvailable()) {
+      } else if (storageType == Type.OSS && location.getStorageInfo.OSSAvailable()) {
         val shuffleDir =
           new Path(new Path(ossDir, conf.workerWorkingDir), s"$appId/$shuffleId")
         FileSystem.mkdirs(