[LLM] adapt pdc sdk

Author: SylarTiaNII

paddlenlp/utils/downloader.py

@@ -38,9 +38,7 @@
     FLASH_DEVICE,
     PDCErrorCode,
     PDCErrorMessageMap,
-    pdc_backup_to_flash_device,
     pdc_flash_device_available,
-    pdc_flash_do_check,
     pdc_tool,
 )
 
@@ -519,6 +517,18 @@ def download_from_pdc(remote_path, local_path, timeout):
 def get_static_model_on_pdc(remote_path, local_path, timeout, enable_flash_device=False):
     """
     Get static model from PDC. Use flash device if possible.
+    This function has to be called after distributed env is initialized in distributed mode.
+    Args:
+        remote_path (`str`):
+            remote path url for download
+        local_path (`str`):
+            local path to place downloaded object
+        timeout (`int`):
+            max wait time for download
+        enable_flash_device (`bool`):
+            Whether to use flash device
+    Returns:
+        str: path to load static model
     """
     try:
         base_dir, target_dir = os.path.split(os.path.normpath(local_path))
@@ -527,6 +537,8 @@ def get_static_model_on_pdc(remote_path, local_path, timeout, enable_flash_devic
     except Exception as e:
         raise RuntimeError(f"{PDC_DOWNLOAD_ERROR}; Failed to parse checkpoint path, details: {e}")
 
+    assert target_dir != ".", f"{PDC_DOWNLOAD_ERROR}, illegal local_path: {local_path}."
+
     flash_path = os.path.join(FLASH_DEVICE, target_dir)
     persistent_path = local_path
 
@@ -543,15 +555,16 @@ def get_static_model_on_pdc(remote_path, local_path, timeout, enable_flash_devic
         logger.info(f"flash device is available, checking status on {flash_path}...")
         # skip download SC as default when flash device is available
         need_download_from_remote = False
-        if os.path.exists(flash_path) and pdc_flash_do_check(flash_path):
+        if os.path.exists(flash_path) and pdc_tool.pdc_flash_do_check(flash_path) == PDCErrorCode.Success:
             logger.info("Static model checked successfully on flash device, ready to load...")
         else:
             logger.warning(
                 "flash device is available but no valid static model found on flash device, need to download from remote."
             )
-            shutil.rmtree(flash_path)
             need_download_from_remote = True
             need_backup_to_flash = True
+    else:
+        logger.info("Flash device is not enabled or available, will download static model from remote.")
 
     # step 2: download from remote if neccesary
     if need_download_from_remote:
@@ -561,12 +574,11 @@ def get_static_model_on_pdc(remote_path, local_path, timeout, enable_flash_devic
 
     # step 3: backup to flash device if flash device is available
     if enable_flash_device and need_backup_to_flash:
-        result = pdc_backup_to_flash_device(persistent_path, flash_path)
+        result = pdc_tool.pdc_backup_to_flash_device(persistent_path, flash_path)
         if result == PDCErrorCode.Success:
             logger.info(f"Backup static model to flash device {flash_path} successfully.")
         else:
             logger.error(f"Backup static model to flash device failed, error details: {PDCErrorMessageMap[result]}.")
-            shutil.rmtree(flash_path)
 
     # step 4: return flash path if available, otherwise return persistent path
     if dist.get_world_size() > 1:

paddlenlp/utils/pdc_sdk.py

@@ -15,9 +15,11 @@
 import json
 import os
 import queue
+import shutil
 import subprocess
 import threading
 import time
+from distutils.dir_util import copy_tree
 from enum import Enum
 from typing import List
 
@@ -28,6 +30,13 @@
 TRAIN_CONFIG = "/root/paddlejob/workspace/env_run/longjob/train.conf"
 TAR_BIN = "tar"
 
+FLASH_DEVICE = os.getenv("PDC_FLASH_DEVICE", "/shared/dev/shm/flash")
+
+
+def pdc_flash_device_available():
+    # TODO(@gexiao): need better check
+    return os.path.exists(FLASH_DEVICE)
+
 
 class PDCErrorCode(Enum):
     """Error Code For PDCTools usage"""
@@ -48,6 +57,7 @@ class PDCErrorCode(Enum):
     InvalidArgument = 1503
     CommandTimeout = 1504
     CheckSumCommandFail = 1505
+    CopyTreeFailed = 1506
 
     UnknownError = 1999
 
@@ -491,14 +501,60 @@ def _download_file(self, remote_path: str, local_path: str) -> PDCErrorCode:
             raise Exception(f"exec cmd {download_cmd_args} with error: {e}")
         return error_code
 
-    def pdc_fc_generate_checksum(self, path: str) -> PDCErrorCode:
+    def _pdc_backup_failed_directory(self, path):
+        base_dir, target_path = os.path.split(os.path.normpath(path))
+        failed_path = os.path.join(base_dir, f"{target_path}_failed")
+        if os.path.exists(path):
+            if os.path.exists(failed_path):
+                shutil.rmtree(failed_path)
+            # Backup failed files for debug
+            os.rename(path, failed_path)
+
+    def pdc_backup_to_flash_device(self, persistent_path: str, flash_device_path: str) -> PDCErrorCode:
+        """backup data to flash device
+
+        Args:
+        persistent_path str: persistent path
+        flash_device_path str: flash device path
+        """
+        if not os.path.exists(persistent_path):
+            logger.error(f"{persistent_path} not exist")
+            return PDCErrorCode.LocalPathNotExist
+
+        logger.info("starting backup to flash device...")
+
+        # step 1: generate checksum for recovery
+        result = self.pdc_generate_dir_checksum(persistent_path)
+        if result != PDCErrorCode.Success:
+            logger.error(f"[Error] [pdc_sdk] generating checksum for {persistent_path} failed")
+            return result
+
+        # step 2: copy persistent data to flash device
+        try:
+            copy_tree(persistent_path, flash_device_path)
+            logger.info(f"backup {persistent_path} to {flash_device_path} successed.")
+        except Exception as e:
+            logger.error(f"[Error] [pdc_sdk] copy tree {persistent_path} to {flash_device_path} failed, error: {e}")
+            self._pdc_backup_failed_directory(flash_device_path)
+            return PDCErrorCode.CopyTreeFailed
+
+        # step 3: do checksum for storage on flash device
+        result = self.pdc_flash_do_check(flash_device_path)
+        if result == PDCErrorCode.Success:
+            return result
+
+        logger.error(f"[Error] [pdc_sdk] checksum failed on {flash_device_path} after copy, backup for debug")
+        self._pdc_backup_failed_directory(flash_device_path)
+        return result
+
+    def pdc_generate_dir_checksum(self, path: str) -> PDCErrorCode:
         """
         Args
         :param localPath:
         :return:
         """
         if not os.path.exists(path):
-            logger.error(f"pdc_fc_generate_checksum gi{path} not exist")
+            logger.error(f"pdc_generate_dir_checksum gi{path} not exist")
             return PDCErrorCode.CommandFail
         generate_checksum_args = [self._pdc_agent_bin, "-mode", "command", "-type", "generateSum", "-path", f"{path}"]
         error_code = PDCErrorCode.Success
@@ -512,14 +568,14 @@ def pdc_fc_generate_checksum(self, path: str) -> PDCErrorCode:
             return PDCErrorCode.CheckSumCommandFail
         return error_code
 
-    def pdc_fc_do_check(self, path: str) -> PDCErrorCode:
+    def pdc_flash_do_check(self, path: str) -> PDCErrorCode:
         """
         Args
         :param localPath:
         :return:
         """
         if not os.path.exists(path):
-            logger.error(f"pdc_fc_do_check {path} not exist")
+            logger.error(f"pdc_flash_do_check {path} not exist")
             return PDCErrorCode.CommandFail
         generate_checksum_args = [self._pdc_agent_bin, "-mode", "command", "-type", "checkSum", "-path", f"{path}"]
         error_code = PDCErrorCode.Success
@@ -528,8 +584,12 @@ def pdc_fc_do_check(self, path: str) -> PDCErrorCode:
             res, error_code = self._exec_cmd(generate_checksum_args)
             if error_code == PDCErrorCode.Success:
                 logger.info(f"check_sum {path} successfully")
+            else:
+                logger.error(f"[Error] [pdc_sdk] check_sum {path} failed, error code: {error_code}")
+                self._pdc_backup_failed_directory(path)
         except Exception as e:
-            logger.error(f"exec cmd {generate_checksum_args} with error: {e}")
+            logger.error(f"[Error] [pdc_sdk] exec cmd {generate_checksum_args} with error: {e}")
+            self._pdc_backup_failed_directory(path)
             return PDCErrorCode.CheckSumCommandFail
         return error_code
 
@@ -562,4 +622,6 @@ def _clean_tmp_files(self, tmp_files: List[str]):
     PDCErrorCode.CalculateHashFail: "calculate hash fail",
     PDCErrorCode.InvalidArgument: "invalid argument",
     PDCErrorCode.CommandTimeout: "command timeout",
+    PDCErrorCode.CheckSumCommandFail: "checksum command fail",
+    PDCErrorCode.CopyTreeFailed: "copy directory failed",
 }