You are viewing a plain text version of this content. The canonical link for it is here.
Posted to reviews@spark.apache.org by GitBox <gi...@apache.org> on 2022/03/18 01:08:00 UTC

[GitHub] [spark] zhouyejoe opened a new pull request #35906: [WIP][SPARK-33236][shuffle] Enable Push-based shuffle service to store state in NM level DB for work preserving restart

zhouyejoe opened a new pull request #35906:
URL: https://github.com/apache/spark/pull/35906


   ### What changes were proposed in this pull request?
   This PR adds the capability of storing the required information into LevelDB for push based shuffle.
   
   ### Why are the changes needed?
   Without this PR, all the information is currently only stored in memory for push based shuffle in shuffle services. During NodeManager restarts, all these information will be lost. Either all the former merged shuffle data won't be able to serve the fetch requests, nor the shuffle services cannot merge any new push blocks from existing applications. After this patch, those information will be stored in LevelDB, and all the information will be recovered during NodeManager restarts.
   
   ### Does this PR introduce _any_ user-facing change?
   No.
   
   ### How was this patch tested?
   Unit testing.
   Deployed to clusters, and restart NMs in 3 different scenarios:
   While a Spark shell application is running, restart NMs in magnet partition in three different scenarios:
   1. After Spark-shell starts, but no scripts running, this will test the NM restart after application/executors register with NMs
   2. While there is on-going shuffle push to shuffle services
   3. While there is on-going merged shuffle fetch from shuffle services.
   Results of the large shuffle testing scripts is identical to the case when there is no NM restart.
   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


[GitHub] [spark] zhouyejoe commented on pull request #35906: [WIP][SPARK-33236][shuffle] Enable Push-based shuffle service to store state in NM level DB for work preserving restart

Posted by GitBox <gi...@apache.org>.
zhouyejoe commented on pull request #35906:
URL: https://github.com/apache/spark/pull/35906#issuecomment-1071921631


   cc @mridulm @otterc 


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


[GitHub] [spark] mridulm commented on a change in pull request #35906: [WIP][SPARK-33236][shuffle] Enable Push-based shuffle service to store state in NM level DB for work preserving restart

Posted by GitBox <gi...@apache.org>.
mridulm commented on a change in pull request #35906:
URL: https://github.com/apache/spark/pull/35906#discussion_r836934153



##########
File path: common/network-shuffle/src/main/java/org/apache/spark/network/shuffle/RemoteBlockPushResolver.java
##########
@@ -655,6 +776,238 @@ public void registerExecutor(String appId, ExecutorShuffleInfo executorInfo) {
     }
   }
 
+
+  @Override
+  public void close() {
+    if (db != null) {
+      try {
+        db.close();
+      } catch (IOException e) {
+        logger.error("Exception closing leveldb with registered app paths info and "
+            + "shuffle partition info", e);
+      }
+    }
+  }
+
+  private static byte[] dbAppAttemptShufflePartitionKeyWithCustomPrefix(
+      AppShufflePartitionId id,
+      String prefix) throws IOException {
+    // we stick a common prefix on all the keys so we can find them in the DB
+    String appAttemptShufflePartitionJson = mapper.writeValueAsString(id);
+    String key = prefix + ";" + appAttemptShufflePartitionJson;
+    return key.getBytes(StandardCharsets.UTF_8);
+  }
+
+  private static AppShufflePartitionId parseDbAppAttemptShufflePartitionKeyWithCustomPrefix(
+      String s,
+      String prefix) throws IOException {
+    if (!s.startsWith(prefix)) {

Review comment:
       nit: check for prefix followed by ';' ?
   Will also prevent potential `ArrayIndexOutOfBoundsException` in `substring` below.

##########
File path: common/network-shuffle/src/main/java/org/apache/spark/network/shuffle/RemoteBlockPushResolver.java
##########
@@ -655,6 +776,238 @@ public void registerExecutor(String appId, ExecutorShuffleInfo executorInfo) {
     }
   }
 
+
+  @Override
+  public void close() {
+    if (db != null) {
+      try {
+        db.close();
+      } catch (IOException e) {
+        logger.error("Exception closing leveldb with registered app paths info and "
+            + "shuffle partition info", e);
+      }
+    }
+  }
+
+  private static byte[] dbAppAttemptShufflePartitionKeyWithCustomPrefix(
+      AppShufflePartitionId id,
+      String prefix) throws IOException {
+    // we stick a common prefix on all the keys so we can find them in the DB
+    String appAttemptShufflePartitionJson = mapper.writeValueAsString(id);
+    String key = prefix + ";" + appAttemptShufflePartitionJson;
+    return key.getBytes(StandardCharsets.UTF_8);
+  }
+
+  private static AppShufflePartitionId parseDbAppAttemptShufflePartitionKeyWithCustomPrefix(
+      String s,
+      String prefix) throws IOException {
+    if (!s.startsWith(prefix)) {
+      throw new IllegalArgumentException("expected a string starting with " + prefix);
+    }
+    String json = s.substring(prefix.length() + 1);
+    return mapper.readValue(json, AppShufflePartitionId.class);
+  }
+
+  private static byte[] dbAppAttemptPathsKey(String appId, int attemptId) throws IOException {
+    // we stick a common prefix on all the keys so we can find them in the DB
+    String appAttemptIdJson = mapper.writeValueAsString(new AppAttemptId(appId, attemptId));
+    String key = APP_ATTEMPT_PATH_KEY_PREFIX + ";" + appAttemptIdJson;
+    return key.getBytes(StandardCharsets.UTF_8);
+  }
+
+  private static AppAttemptId parseDbAppAttemptPathsKey(String s) throws IOException {
+    if (!s.startsWith(APP_ATTEMPT_PATH_KEY_PREFIX)) {
+      throw new IllegalArgumentException("expected a string starting with "
+        + APP_ATTEMPT_PATH_KEY_PREFIX);
+    }
+    String json = s.substring(APP_ATTEMPT_PATH_KEY_PREFIX.length() + 1);
+    return mapper.readValue(json, AppAttemptId.class);
+  }
+
+  public void reloadAppShuffleInfo(DB db) throws IOException {
+    logger.debug("Reload applications merged shuffle information from DB");
+    reloadActiveAppAttemptsPathInfo(db);
+    reloadFinalizedAppAttemptsShuffleInfo(db);
+    reloadActiveAppAttemptsShufflePartitionInfo(db);

Review comment:
       `reloadAppShuffleInfo` is used for testing, not the others.
   Make `reloadAppShuffleInfo` package private with `@VisibleForTesting` and make other methods `private` ?
   If we need to open these up in future for testing, we can do so at that time.

##########
File path: common/network-shuffle/src/main/java/org/apache/spark/network/shuffle/RemoteBlockPushResolver.java
##########
@@ -655,6 +776,238 @@ public void registerExecutor(String appId, ExecutorShuffleInfo executorInfo) {
     }
   }
 
+
+  @Override
+  public void close() {
+    if (db != null) {
+      try {
+        db.close();
+      } catch (IOException e) {
+        logger.error("Exception closing leveldb with registered app paths info and "
+            + "shuffle partition info", e);
+      }
+    }
+  }
+
+  private static byte[] dbAppAttemptShufflePartitionKeyWithCustomPrefix(
+      AppShufflePartitionId id,
+      String prefix) throws IOException {
+    // we stick a common prefix on all the keys so we can find them in the DB
+    String appAttemptShufflePartitionJson = mapper.writeValueAsString(id);
+    String key = prefix + ";" + appAttemptShufflePartitionJson;
+    return key.getBytes(StandardCharsets.UTF_8);
+  }
+
+  private static AppShufflePartitionId parseDbAppAttemptShufflePartitionKeyWithCustomPrefix(
+      String s,
+      String prefix) throws IOException {
+    if (!s.startsWith(prefix)) {
+      throw new IllegalArgumentException("expected a string starting with " + prefix);
+    }
+    String json = s.substring(prefix.length() + 1);
+    return mapper.readValue(json, AppShufflePartitionId.class);
+  }
+
+  private static byte[] dbAppAttemptPathsKey(String appId, int attemptId) throws IOException {
+    // we stick a common prefix on all the keys so we can find them in the DB
+    String appAttemptIdJson = mapper.writeValueAsString(new AppAttemptId(appId, attemptId));
+    String key = APP_ATTEMPT_PATH_KEY_PREFIX + ";" + appAttemptIdJson;
+    return key.getBytes(StandardCharsets.UTF_8);
+  }
+
+  private static AppAttemptId parseDbAppAttemptPathsKey(String s) throws IOException {
+    if (!s.startsWith(APP_ATTEMPT_PATH_KEY_PREFIX)) {
+      throw new IllegalArgumentException("expected a string starting with "
+        + APP_ATTEMPT_PATH_KEY_PREFIX);
+    }
+    String json = s.substring(APP_ATTEMPT_PATH_KEY_PREFIX.length() + 1);
+    return mapper.readValue(json, AppAttemptId.class);
+  }
+
+  public void reloadAppShuffleInfo(DB db) throws IOException {
+    logger.debug("Reload applications merged shuffle information from DB");
+    reloadActiveAppAttemptsPathInfo(db);
+    reloadFinalizedAppAttemptsShuffleInfo(db);
+    reloadActiveAppAttemptsShufflePartitionInfo(db);
+  }
+
+  @VisibleForTesting
+  public void reloadActiveAppAttemptsPathInfo(DB db) throws IOException {
+    if (db != null) {
+      DBIterator itr = db.iterator();
+      itr.seek(APP_ATTEMPT_PATH_KEY_PREFIX.getBytes(StandardCharsets.UTF_8));
+      while (itr.hasNext()) {
+        Map.Entry<byte[], byte[]> e = itr.next();
+        String key = new String(e.getKey(), StandardCharsets.UTF_8);
+        if (!key.startsWith(APP_ATTEMPT_PATH_KEY_PREFIX)) {
+          break;
+        }
+        AppAttemptId appAttemptId = parseDbAppAttemptPathsKey(key);
+        try{
+          AppPathsInfo appPathsInfo = mapper.readValue(e.getValue(), AppPathsInfo.class);
+          logger.debug("Reloading active application {}_{} merged shuffle files paths",
+            appAttemptId.appId, appAttemptId.attemptId);
+          appsShuffleInfo.computeIfAbsent(appAttemptId.appId, id ->
+            new AppShuffleInfo(appAttemptId.appId, appAttemptId.attemptId, appPathsInfo));
+        } catch (Exception exception) {
+          logger.error("Parsing exception is {}", exception);
+        }
+      }
+    }
+  }
+
+  @VisibleForTesting
+  public void reloadFinalizedAppAttemptsShuffleInfo(DB db) throws IOException {
+    if (db != null) {
+      DBIterator itr = db.iterator();
+      itr.seek(APP_ATTEMPT_SHUFFLE_FINALIZE_STATUS_KEY_PREFIX.getBytes(StandardCharsets.UTF_8));
+      while (itr.hasNext()) {
+        Map.Entry<byte[], byte[]> e = itr.next();
+        String key = new String(e.getKey(), StandardCharsets.UTF_8);
+        if (!key.startsWith(APP_ATTEMPT_SHUFFLE_FINALIZE_STATUS_KEY_PREFIX)) {
+          break;
+        }
+        AppShufflePartitionId partitionId =
+          parseDbAppAttemptShufflePartitionKeyWithCustomPrefix(
+            key, APP_ATTEMPT_SHUFFLE_FINALIZE_STATUS_KEY_PREFIX);
+        if (partitionId.reduceId != -1) {

Review comment:
       Why would reducerId not be -1 for `APP_ATTEMPT_SHUFFLE_FINALIZE_STATUS_KEY_PREFIX` ?

##########
File path: common/network-shuffle/src/main/java/org/apache/spark/network/shuffle/RemoteBlockPushResolver.java
##########
@@ -536,9 +645,20 @@ public MergeStatuses finalizeShuffleMerge(FinalizeShuffleMerge msg) {
         }
       }
       // Even when the mergePartitionsInfo is null, we mark the shuffle as finalized but the results
-      // sent to the driver will be empty. This cam happen when the service didn't receive any
+      // sent to the driver will be empty. This can happen when the service didn't receive any
       // blocks for the shuffle yet and the driver didn't wait for enough time to finalize the
       // shuffle.
+      if (db != null) {

Review comment:
       Do this only for the last `else` block above.

##########
File path: common/network-shuffle/src/main/java/org/apache/spark/network/shuffle/RemoteBlockPushResolver.java
##########
@@ -342,6 +404,30 @@ void closeAndDeletePartitionFilesIfNeeded(
     if (cleanupLocalDirs) {
       deleteExecutorDirs(appShuffleInfo);
     }
+    cleanUpAppShuffleInfoInDB(appShuffleInfo);
+  }
+
+  void cleanUpAppShuffleInfoInDB(AppShuffleInfo appShuffleInfo) {

Review comment:
       nit: limit visibility to private if we dont need to expose it for testing/functionality.

##########
File path: common/network-shuffle/src/main/java/org/apache/spark/network/shuffle/RemoteBlockPushResolver.java
##########
@@ -342,6 +404,30 @@ void closeAndDeletePartitionFilesIfNeeded(
     if (cleanupLocalDirs) {
       deleteExecutorDirs(appShuffleInfo);
     }
+    cleanUpAppShuffleInfoInDB(appShuffleInfo);
+  }
+
+  void cleanUpAppShuffleInfoInDB(AppShuffleInfo appShuffleInfo) {
+    if (db != null) {
+      try {
+        db.delete(dbAppAttemptPathsKey(appShuffleInfo.appId, appShuffleInfo.attemptId));
+        appShuffleInfo.shuffles
+          .forEach((shuffleId, shuffleInfo) -> shuffleInfo.shuffleMergePartitions
+            .forEach((shuffleMergeId, partitionInfo) -> {
+              synchronized (partitionInfo) {
+                if (shuffleInfo.isFinalized()) {
+                  cleanUpAppShufflePartitionInfoInDB(new AppShufflePartitionId(
+                    appShuffleInfo.appId, appShuffleInfo.attemptId, shuffleId,
+                      shuffleMergeId, -1));
+                }

Review comment:
       nit: this `if isFinalized` can be moved into `closeAndDeletePartitionFilesIfNeeded` and we can remove `cleanUpAppShuffleInfoInDB` ?

##########
File path: common/network-shuffle/src/main/java/org/apache/spark/network/shuffle/RemoteBlockPushResolver.java
##########
@@ -655,6 +776,238 @@ public void registerExecutor(String appId, ExecutorShuffleInfo executorInfo) {
     }
   }
 
+
+  @Override
+  public void close() {
+    if (db != null) {
+      try {
+        db.close();
+      } catch (IOException e) {
+        logger.error("Exception closing leveldb with registered app paths info and "
+            + "shuffle partition info", e);
+      }
+    }
+  }
+
+  private static byte[] dbAppAttemptShufflePartitionKeyWithCustomPrefix(
+      AppShufflePartitionId id,
+      String prefix) throws IOException {
+    // we stick a common prefix on all the keys so we can find them in the DB
+    String appAttemptShufflePartitionJson = mapper.writeValueAsString(id);
+    String key = prefix + ";" + appAttemptShufflePartitionJson;
+    return key.getBytes(StandardCharsets.UTF_8);
+  }
+
+  private static AppShufflePartitionId parseDbAppAttemptShufflePartitionKeyWithCustomPrefix(
+      String s,

Review comment:
       nit: `s` -> `key` ?

##########
File path: common/network-shuffle/src/main/java/org/apache/spark/network/shuffle/RemoteBlockPushResolver.java
##########
@@ -261,7 +321,8 @@ public MergedBlockMeta getMergedBlockMeta(
     int size = (int) indexFile.length();
     // First entry is the zero offset
     int numChunks = (size / Long.BYTES) - 1;
-    File metaFile = appShuffleInfo.getMergedShuffleMetaFile(shuffleId, shuffleMergeId, reduceId);
+    File metaFile =
+      appShuffleInfo.getMergedShuffleMetaFile(shuffleId, shuffleMergeId, reduceId);

Review comment:
       revert this and other whitespace changes to minimize diff ?

##########
File path: common/network-shuffle/src/main/java/org/apache/spark/network/shuffle/RemoteBlockPushResolver.java
##########
@@ -617,8 +738,8 @@ public void registerExecutor(String appId, ExecutorShuffleInfo executorInfo) {
           appsShuffleInfo.computeIfAbsent(appId, id ->
             new AppShuffleInfo(
               appId, UNDEFINED_ATTEMPT_ID,
-              new AppPathsInfo(appId, executorInfo.localDirs,
-                mergeDir, executorInfo.subDirsPerLocalDir)
+              new AppPathsInfo(appId, attemptId, executorInfo.localDirs,
+                mergeDir, executorInfo.subDirsPerLocalDir, db)

Review comment:
       Move the db update out of the constructor, and do it explicitly ?

##########
File path: common/network-yarn/src/main/java/org/apache/spark/network/yarn/YarnShuffleService.java
##########
@@ -120,6 +120,7 @@
 
   private static final String RECOVERY_FILE_NAME = "registeredExecutors.ldb";
   private static final String SECRETS_RECOVERY_FILE_NAME = "sparkShuffleRecovery.ldb";
+  private static final String MERGE_MANAGER_FILE_NAME = "mergeManager.ldb";

Review comment:
       nit: rename to `sparkShuffleMergeRecovery` or some such to be more clear about the intent of the file ?

##########
File path: common/network-shuffle/src/main/java/org/apache/spark/network/shuffle/RemoteBlockPushResolver.java
##########
@@ -655,6 +776,238 @@ public void registerExecutor(String appId, ExecutorShuffleInfo executorInfo) {
     }
   }
 
+
+  @Override
+  public void close() {
+    if (db != null) {
+      try {
+        db.close();
+      } catch (IOException e) {
+        logger.error("Exception closing leveldb with registered app paths info and "
+            + "shuffle partition info", e);
+      }
+    }
+  }
+
+  private static byte[] dbAppAttemptShufflePartitionKeyWithCustomPrefix(
+      AppShufflePartitionId id,
+      String prefix) throws IOException {
+    // we stick a common prefix on all the keys so we can find them in the DB
+    String appAttemptShufflePartitionJson = mapper.writeValueAsString(id);
+    String key = prefix + ";" + appAttemptShufflePartitionJson;
+    return key.getBytes(StandardCharsets.UTF_8);
+  }
+
+  private static AppShufflePartitionId parseDbAppAttemptShufflePartitionKeyWithCustomPrefix(
+      String s,
+      String prefix) throws IOException {
+    if (!s.startsWith(prefix)) {
+      throw new IllegalArgumentException("expected a string starting with " + prefix);
+    }
+    String json = s.substring(prefix.length() + 1);
+    return mapper.readValue(json, AppShufflePartitionId.class);
+  }
+
+  private static byte[] dbAppAttemptPathsKey(String appId, int attemptId) throws IOException {
+    // we stick a common prefix on all the keys so we can find them in the DB
+    String appAttemptIdJson = mapper.writeValueAsString(new AppAttemptId(appId, attemptId));
+    String key = APP_ATTEMPT_PATH_KEY_PREFIX + ";" + appAttemptIdJson;
+    return key.getBytes(StandardCharsets.UTF_8);
+  }
+
+  private static AppAttemptId parseDbAppAttemptPathsKey(String s) throws IOException {
+    if (!s.startsWith(APP_ATTEMPT_PATH_KEY_PREFIX)) {
+      throw new IllegalArgumentException("expected a string starting with "
+        + APP_ATTEMPT_PATH_KEY_PREFIX);
+    }
+    String json = s.substring(APP_ATTEMPT_PATH_KEY_PREFIX.length() + 1);
+    return mapper.readValue(json, AppAttemptId.class);
+  }
+
+  public void reloadAppShuffleInfo(DB db) throws IOException {
+    logger.debug("Reload applications merged shuffle information from DB");
+    reloadActiveAppAttemptsPathInfo(db);
+    reloadFinalizedAppAttemptsShuffleInfo(db);
+    reloadActiveAppAttemptsShufflePartitionInfo(db);
+  }
+
+  @VisibleForTesting
+  public void reloadActiveAppAttemptsPathInfo(DB db) throws IOException {
+    if (db != null) {
+      DBIterator itr = db.iterator();
+      itr.seek(APP_ATTEMPT_PATH_KEY_PREFIX.getBytes(StandardCharsets.UTF_8));
+      while (itr.hasNext()) {
+        Map.Entry<byte[], byte[]> e = itr.next();
+        String key = new String(e.getKey(), StandardCharsets.UTF_8);
+        if (!key.startsWith(APP_ATTEMPT_PATH_KEY_PREFIX)) {
+          break;
+        }
+        AppAttemptId appAttemptId = parseDbAppAttemptPathsKey(key);
+        try{
+          AppPathsInfo appPathsInfo = mapper.readValue(e.getValue(), AppPathsInfo.class);
+          logger.debug("Reloading active application {}_{} merged shuffle files paths",
+            appAttemptId.appId, appAttemptId.attemptId);
+          appsShuffleInfo.computeIfAbsent(appAttemptId.appId, id ->
+            new AppShuffleInfo(appAttemptId.appId, appAttemptId.attemptId, appPathsInfo));

Review comment:
       Instead of `computeIfAbsent`, use `compute` and ensure that the newer attempt id gets added (in case an earlier deletion did failed/etc).
   
   Same for `reloadFinalizedAppAttemptsShuffleInfo` below as well (w.r.t `shuffleMergeId`: there can be a race between `closeAndDeletePartitionFiles`, `finalizeShuffleMerge` and NM shutdown).
   
   I am assuming `reloadActiveAppAttemptsShufflePartitionInfo`, `reloadPartitionInfo`, etc will get deleted if we are only tracking finalized shuffles.

##########
File path: common/network-shuffle/src/main/java/org/apache/spark/network/shuffle/RemoteBlockPushResolver.java
##########
@@ -210,8 +252,23 @@ private AppShufflePartitionInfo getOrCreateAppShufflePartitionInfo(
       File metaFile =
         appShuffleInfo.getMergedShuffleMetaFile(shuffleId, shuffleMergeId, reduceId);
       try {
-        return newAppShufflePartitionInfo(appShuffleInfo.appId, shuffleId, shuffleMergeId,
-          reduceId, dataFile, indexFile, metaFile);
+        AppShufflePartitionInfo appShufflePartitionInfo =
+            newAppShufflePartitionInfo(appShuffleInfo.appId, appShuffleInfo.attemptId,
+            shuffleId, shuffleMergeId, reduceId, dataFile, indexFile, metaFile);
+        if (db != null) {

Review comment:
       Why not track only those which have been successfully finalized ?
   That is, we dont track in progress pushes ... this should simply the code, remove all changes to `AppShufflePartitionInfo`, remove `APP_ATTEMPT_SHUFFLE_PARTITION_KEY_PREFIX`, recovery,  etc.
   
   (Some of my other comments in this PR might not be relevant once this is done btw).

##########
File path: common/network-shuffle/src/main/java/org/apache/spark/network/shuffle/RemoteBlockPushResolver.java
##########
@@ -210,8 +252,23 @@ private AppShufflePartitionInfo getOrCreateAppShufflePartitionInfo(
       File metaFile =
         appShuffleInfo.getMergedShuffleMetaFile(shuffleId, shuffleMergeId, reduceId);
       try {
-        return newAppShufflePartitionInfo(appShuffleInfo.appId, shuffleId, shuffleMergeId,
-          reduceId, dataFile, indexFile, metaFile);
+        AppShufflePartitionInfo appShufflePartitionInfo =
+            newAppShufflePartitionInfo(appShuffleInfo.appId, appShuffleInfo.attemptId,
+            shuffleId, shuffleMergeId, reduceId, dataFile, indexFile, metaFile);
+        if (db != null) {
+          try{
+            byte[] dbKey =
+                dbAppAttemptShufflePartitionKeyWithCustomPrefix(

Review comment:
       nit: Add `getShufflePartitionDbKey` to be more clear ? (which delegates to `dbAppAttemptShufflePartitionKeyWithCustomPrefix` if required). The read/parse logic seems to be doing it with `parseDbApp*` methods already.
   
   Same for other uses of `dbAppAttemptShufflePartitionKeyWithCustomPrefix` as well.
   Localize all string/serde manipulations to a small set of methods
   

##########
File path: common/network-shuffle/src/main/java/org/apache/spark/network/shuffle/RemoteBlockPushResolver.java
##########
@@ -655,6 +776,238 @@ public void registerExecutor(String appId, ExecutorShuffleInfo executorInfo) {
     }
   }
 
+
+  @Override
+  public void close() {
+    if (db != null) {
+      try {
+        db.close();
+      } catch (IOException e) {
+        logger.error("Exception closing leveldb with registered app paths info and "
+            + "shuffle partition info", e);
+      }
+    }
+  }
+
+  private static byte[] dbAppAttemptShufflePartitionKeyWithCustomPrefix(
+      AppShufflePartitionId id,
+      String prefix) throws IOException {
+    // we stick a common prefix on all the keys so we can find them in the DB
+    String appAttemptShufflePartitionJson = mapper.writeValueAsString(id);
+    String key = prefix + ";" + appAttemptShufflePartitionJson;

Review comment:
       nit: pull the `;` as a constant

##########
File path: common/network-shuffle/src/main/java/org/apache/spark/network/shuffle/RemoteBlockPushResolver.java
##########
@@ -210,8 +252,23 @@ private AppShufflePartitionInfo getOrCreateAppShufflePartitionInfo(
       File metaFile =
         appShuffleInfo.getMergedShuffleMetaFile(shuffleId, shuffleMergeId, reduceId);
       try {
-        return newAppShufflePartitionInfo(appShuffleInfo.appId, shuffleId, shuffleMergeId,
-          reduceId, dataFile, indexFile, metaFile);
+        AppShufflePartitionInfo appShufflePartitionInfo =
+            newAppShufflePartitionInfo(appShuffleInfo.appId, appShuffleInfo.attemptId,
+            shuffleId, shuffleMergeId, reduceId, dataFile, indexFile, metaFile);
+        if (db != null) {
+          try{
+            byte[] dbKey =
+                dbAppAttemptShufflePartitionKeyWithCustomPrefix(
+                  appShufflePartitionInfo.appShufflePartitionId,
+                  APP_ATTEMPT_SHUFFLE_PARTITION_KEY_PREFIX);
+            // We can reload a partition's state solely from the partition ID. Thus only the
+            // partition Ids are stored in levelDB.
+            db.put(dbKey, new byte[0]);
+          } catch (Exception e) {
+            logger.error("Error saving active app shuffle partition", e);
+          }
+        }

Review comment:
       We have a bunch of 
   ```
   if (db != null) { 
     try {
       byte[] key = generate key
       db.put(key, value)
     } catch (Exception ex) {
       log message
     }
   }
   ```
   sprinkled across this class.
   
   Refactor these and move them to util methods ? So that the main flow is clearer ?
   
   ```suggestion
     writeShufflePartitionToDb(appShufflePartitionInfo.appShufflePartitionId)
   ```

##########
File path: common/network-shuffle/src/main/java/org/apache/spark/network/shuffle/RemoteBlockPushResolver.java
##########
@@ -355,8 +441,31 @@ void closeAndDeletePartitionFiles(Map<Integer, AppShufflePartitionInfo> partitio
       .forEach((partitionId, partitionInfo) -> {
         synchronized (partitionInfo) {
           partitionInfo.closeAllFilesAndDeleteIfNeeded(true);
+          cleanUpAppShufflePartitionInfoInDB(partitionInfo.appShufflePartitionId);
         }
       });
+
+  }
+
+  void cleanUpAppShufflePartitionInfoInDB(AppShufflePartitionId appShufflePartitionId) {
+    if (db != null) {
+      try {
+        if (appShufflePartitionId.reduceId == -1) {
+          db.delete(dbAppAttemptShufflePartitionKeyWithCustomPrefix(
+            appShufflePartitionId, APP_ATTEMPT_SHUFFLE_FINALIZE_STATUS_KEY_PREFIX));
+        } else {
+          db.delete(dbAppAttemptShufflePartitionKeyWithCustomPrefix(
+            appShufflePartitionId, APP_ATTEMPT_SHUFFLE_PARTITION_KEY_PREFIX));
+        }
+      } catch (IOException e) {
+        logger.error("Error deleting {}_{} shuffleMergeId {} shuffle {}"
+                + "reduce {} from application shuffle merged partition info db",
+            appShufflePartitionId.appId, appShufflePartitionId.attemptId,
+            appShufflePartitionId.shuffleMergeId,
+            appShufflePartitionId.shuffleId,

Review comment:
       `shuffleId` before `shuffleMergeId` (here and in `updateChunkInfo` - not sure if elsewhere as well).
   Also, use `toString` in `AppShufflePartitionId` instead (here and elsewhere we are logging) ?

##########
File path: common/network-yarn/src/main/java/org/apache/spark/network/yarn/YarnShuffleService.java
##########
@@ -296,7 +307,8 @@ static MergedShuffleFileManager newMergedShuffleFileManagerInstance(TransportCon
         mergeManagerImplClazz.asSubclass(MergedShuffleFileManager.class);
       // The assumption is that all the custom implementations just like the RemoteBlockPushResolver
       // will also need the transport configuration.
-      return mergeManagerSubClazz.getConstructor(TransportConf.class).newInstance(conf);
+      return mergeManagerSubClazz.getConstructor(TransportConf.class, File.class)
+        .newInstance(conf, mergeManagerFile);

Review comment:
       Modify `NoOpMergedShuffleFileManager` to also support this constructor - check other implementations as well if relevant.

##########
File path: common/network-shuffle/src/main/java/org/apache/spark/network/shuffle/RemoteBlockPushResolver.java
##########
@@ -655,6 +776,238 @@ public void registerExecutor(String appId, ExecutorShuffleInfo executorInfo) {
     }
   }
 
+
+  @Override
+  public void close() {
+    if (db != null) {
+      try {
+        db.close();
+      } catch (IOException e) {
+        logger.error("Exception closing leveldb with registered app paths info and "
+            + "shuffle partition info", e);
+      }
+    }
+  }
+
+  private static byte[] dbAppAttemptShufflePartitionKeyWithCustomPrefix(
+      AppShufflePartitionId id,
+      String prefix) throws IOException {
+    // we stick a common prefix on all the keys so we can find them in the DB
+    String appAttemptShufflePartitionJson = mapper.writeValueAsString(id);
+    String key = prefix + ";" + appAttemptShufflePartitionJson;
+    return key.getBytes(StandardCharsets.UTF_8);
+  }
+
+  private static AppShufflePartitionId parseDbAppAttemptShufflePartitionKeyWithCustomPrefix(
+      String s,
+      String prefix) throws IOException {
+    if (!s.startsWith(prefix)) {
+      throw new IllegalArgumentException("expected a string starting with " + prefix);
+    }
+    String json = s.substring(prefix.length() + 1);
+    return mapper.readValue(json, AppShufflePartitionId.class);
+  }
+
+  private static byte[] dbAppAttemptPathsKey(String appId, int attemptId) throws IOException {
+    // we stick a common prefix on all the keys so we can find them in the DB
+    String appAttemptIdJson = mapper.writeValueAsString(new AppAttemptId(appId, attemptId));
+    String key = APP_ATTEMPT_PATH_KEY_PREFIX + ";" + appAttemptIdJson;
+    return key.getBytes(StandardCharsets.UTF_8);
+  }
+
+  private static AppAttemptId parseDbAppAttemptPathsKey(String s) throws IOException {
+    if (!s.startsWith(APP_ATTEMPT_PATH_KEY_PREFIX)) {
+      throw new IllegalArgumentException("expected a string starting with "
+        + APP_ATTEMPT_PATH_KEY_PREFIX);
+    }
+    String json = s.substring(APP_ATTEMPT_PATH_KEY_PREFIX.length() + 1);
+    return mapper.readValue(json, AppAttemptId.class);
+  }
+
+  public void reloadAppShuffleInfo(DB db) throws IOException {
+    logger.debug("Reload applications merged shuffle information from DB");
+    reloadActiveAppAttemptsPathInfo(db);
+    reloadFinalizedAppAttemptsShuffleInfo(db);
+    reloadActiveAppAttemptsShufflePartitionInfo(db);
+  }
+
+  @VisibleForTesting
+  public void reloadActiveAppAttemptsPathInfo(DB db) throws IOException {
+    if (db != null) {
+      DBIterator itr = db.iterator();
+      itr.seek(APP_ATTEMPT_PATH_KEY_PREFIX.getBytes(StandardCharsets.UTF_8));
+      while (itr.hasNext()) {
+        Map.Entry<byte[], byte[]> e = itr.next();
+        String key = new String(e.getKey(), StandardCharsets.UTF_8);
+        if (!key.startsWith(APP_ATTEMPT_PATH_KEY_PREFIX)) {
+          break;
+        }
+        AppAttemptId appAttemptId = parseDbAppAttemptPathsKey(key);
+        try{
+          AppPathsInfo appPathsInfo = mapper.readValue(e.getValue(), AppPathsInfo.class);
+          logger.debug("Reloading active application {}_{} merged shuffle files paths",
+            appAttemptId.appId, appAttemptId.attemptId);
+          appsShuffleInfo.computeIfAbsent(appAttemptId.appId, id ->
+            new AppShuffleInfo(appAttemptId.appId, appAttemptId.attemptId, appPathsInfo));
+        } catch (Exception exception) {
+          logger.error("Parsing exception is {}", exception);
+        }
+      }
+    }
+  }
+
+  @VisibleForTesting
+  public void reloadFinalizedAppAttemptsShuffleInfo(DB db) throws IOException {
+    if (db != null) {
+      DBIterator itr = db.iterator();
+      itr.seek(APP_ATTEMPT_SHUFFLE_FINALIZE_STATUS_KEY_PREFIX.getBytes(StandardCharsets.UTF_8));
+      while (itr.hasNext()) {
+        Map.Entry<byte[], byte[]> e = itr.next();
+        String key = new String(e.getKey(), StandardCharsets.UTF_8);
+        if (!key.startsWith(APP_ATTEMPT_SHUFFLE_FINALIZE_STATUS_KEY_PREFIX)) {
+          break;
+        }
+        AppShufflePartitionId partitionId =
+          parseDbAppAttemptShufflePartitionKeyWithCustomPrefix(
+            key, APP_ATTEMPT_SHUFFLE_FINALIZE_STATUS_KEY_PREFIX);
+        if (partitionId.reduceId != -1) {
+          break;
+        }
+        AppShuffleInfo appShuffleInfo = appsShuffleInfo.get(partitionId.appId);
+        if (appShuffleInfo != null) {
+          appShuffleInfo.shuffles.putIfAbsent(partitionId.shuffleId,
+            new AppShuffleMergePartitionsInfo(partitionId.shuffleMergeId, true));

Review comment:
       See comment in `reloadActiveAppAttemptsPathInfo` above.

##########
File path: common/network-shuffle/src/main/java/org/apache/spark/network/shuffle/RemoteBlockPushResolver.java
##########
@@ -1172,47 +1634,58 @@ private void finalizePartition() throws IOException {
       // Get rid of any partial block data at the end of the file. This could either
       // be due to failure, or a request still being processed when the shuffle
       // merge gets finalized, or any exceptions while updating index/meta files.
-      dataChannel.truncate(lastChunkOffset);
+      logger.trace("{}_{} shuffleId {} shuffleMergeId {} reduceId {} "
+         + "truncating files data {} index {} meta {}", appShufflePartitionId.appId,
+          appShufflePartitionId.attemptId, appShufflePartitionId.shuffleId,
+          appShufflePartitionId.shuffleMergeId, appShufflePartitionId.reduceId,
+          lastChunkOffset, indexFile.getPos(), metaFile.getPos());
+      dataFile.getChannel().truncate(lastChunkOffset);
       indexFile.getChannel().truncate(indexFile.getPos());
       metaFile.getChannel().truncate(metaFile.getPos());
     }
 
     void closeAllFilesAndDeleteIfNeeded(boolean delete) {
       try {
-        if (dataChannel.isOpen()) {
-          dataChannel.close();
-          if (delete) {
-            dataFile.delete();
-          }
+        dataFile.close();
+        if (delete) {
+          dataFile.delete();

Review comment:
       There is a change in behavior here - we are deleting even if file was not open.
   If we are simplifying to only tracking finalized merges, we remove the need for this change as well and preserve earlier behavior.




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


[GitHub] [spark] zhouyejoe commented on pull request #35906: [WIP][SPARK-33236][shuffle] Enable Push-based shuffle service to store state in NM level DB for work preserving restart

Posted by GitBox <gi...@apache.org>.
zhouyejoe commented on pull request #35906:
URL: https://github.com/apache/spark/pull/35906#issuecomment-1071921631


   cc @mridulm @otterc 


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


[GitHub] [spark] AmplabJenkins commented on pull request #35906: [WIP][SPARK-33236][shuffle] Enable Push-based shuffle service to store state in NM level DB for work preserving restart

Posted by GitBox <gi...@apache.org>.
AmplabJenkins commented on pull request #35906:
URL: https://github.com/apache/spark/pull/35906#issuecomment-1073022301


   Can one of the admins verify this patch?


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org