Files
wehub-resource-sync 2aaeece67c
Codestyle Check / Lint (push) Has been cancelled
Codestyle Check / Check bypass (push) Has been cancelled
Pipelines-Test / Pipelines-Test (push) Has been cancelled
chore: import upstream snapshot with attribution
2026-07-13 13:37:14 +08:00

630 lines
24 KiB
Python

# Copyright (c) 2024 PaddlePaddle Authors. 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.
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
from paddlenlp.utils.log import logger
PDC_AGENT_BIN = "/root/paddlejob/tools/agent"
HASH_SUM_BIN = "/root/paddlejob/afs_tool/bin/b3sum"
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"""
# success
Success = 0
RemotePathNotExist = 1404
LocalPathExist = 1405
DownloadFail = 1406
AgentConfigInvalid = 1407
AFSToolsNotExist = 1408
TrainConfigNotExist = 1409
LocalPathNotExist = 1410
CommandFail = 1501
CalculateHashFail = 1502
InvalidArgument = 1503
CommandTimeout = 1504
CheckSumCommandFail = 1505
CopyTreeFailed = 1506
UnknownError = 1999
class PDCTools:
"""PDCTools"""
def __init__(self):
""" """
self._pdc_agent_bin = PDC_AGENT_BIN
self._hash_sum_bin = HASH_SUM_BIN
self._train_config = TRAIN_CONFIG
self._tar_bin = TAR_BIN
def pdc_upload(self, remote_path: str, local_path: str) -> PDCErrorCode:
"""upload data to afs/bos
1. tar local path
2. calculate the hash of tar file
3. upload tar to remote path
Args:
remote_path str: the remote file path, afs/bos, such as afs://user/a/b/xx.tar
local_path str: local file path
Return:
"""
pre_check_status = self._pre_check()
if pre_check_status != PDCErrorCode.Success:
return pre_check_status
# check local path
if not os.path.exists(local_path):
logger.error(f"{local_path} not exist")
return PDCErrorCode.LocalPathNotExist
if not remote_path.endswith(".tar"):
logger.warning(f"remote path {remote_path} should end with .tar")
return PDCErrorCode.InvalidArgument
try:
# get tar name
remote_dir = os.path.dirname(remote_path)
tar_file_name = os.path.basename(remote_path)
# tar local path
status = self._tar_file(local_path, tar_file_name)
if status != PDCErrorCode.Success:
logger.error(f"tar local path {local_path} failed")
return status
# calc hash
b3sum_hash, status = self._calculate_hash(tar_file_name)
if status != PDCErrorCode.Success:
logger.error(f"calculate hash for {tar_file_name} failed")
return status
logger.info(f"local tar: {tar_file_name}, b3sum hash: {b3sum_hash}")
# upload local tar to remote path
status = self._upload_file(tar_file_name, remote_path)
if status != PDCErrorCode.Success:
logger.error(f"upload file {tar_file_name} failed")
return status
# upload b3sum hash to remote path
local_b3sum_hash_file = f".{time.time()}_b3sum.hash"
with open(local_b3sum_hash_file, "w") as f:
f.write(b3sum_hash)
remote_b3sum_path = os.path.join(remote_dir, self._get_file_hash_name(tar_file_name))
status = self._upload_file(local_b3sum_hash_file, remote_b3sum_path)
if status != PDCErrorCode.Success:
logger.error(f"upload hash file {local_b3sum_hash_file} failed")
return status
# clean tmp files
self._clean_tmp_files([local_b3sum_hash_file, tar_file_name])
logger.info(f"successfully uploaded ${local_path} to remote path ${remote_path}")
return PDCErrorCode.Success
except Exception as e:
logger.error(f"pdc upload failed: {e}")
raise e
def pdc_download(self, remote_path: str, local_path: str, timeout: int) -> PDCErrorCode:
"""
download data from afs/bos
Args:
remote_path str: the remote file path, afs/bos, such as afs://user/a/b/xx.tar
local_path str: local file directory
timeout int: max wait time in seconds
Return:
PDCErrorCode: indicate the status of pdc download
"""
def _download_worker(remote_path, local_path, queue):
try:
result = self._pdc_download_impl(remote_path, local_path)
queue.put(result)
except Exception as e:
queue.put(str(e))
return
result_queue = queue.Queue()
thread = threading.Thread(
target=_download_worker,
args=(
remote_path,
local_path,
result_queue,
),
)
thread.start()
logger.info(f"Begin downloading object of {remote_path} to {local_path} from PDC...")
start_time = time.time()
end_time = time.time()
last_log_time = start_time
while (end_time - start_time) < timeout:
if not thread.is_alive():
break
if end_time - last_log_time > 30:
# log every 30 seconds to avoid false detection by hangWatcher
logger.info(f"Still waiting for download, already passed {end_time - start_time} seconds...")
last_log_time = end_time
time.sleep(1)
end_time = time.time()
if thread.is_alive():
return PDCErrorCode.CommandTimeout
result = result_queue.get()
if isinstance(result, str):
logger.error(f"Unknown exception occurred during download process, details: {result}")
return PDCErrorCode.UnknownError
return result
def pdc_download_checkpoint(self, resume_step: int, timeout: int) -> PDCErrorCode:
"""
download checkpoints from afs/bos
Args:
resume_step int: the resume step number
timeout int: max wait time in seconds
Return:
PDCErrorCode: indicate the status of pdc download
"""
def _download_worker(step, queue):
try:
result = self._pdc_download_checkpoint_impl(step)
queue.put(result)
except Exception as e:
queue.put(str(e))
return
result_queue = queue.Queue()
thread = threading.Thread(
target=_download_worker,
args=(
resume_step,
result_queue,
),
)
thread.start()
logger.info(f"Begin downloading recovery checkpoint of step {resume_step} from PDC...")
start_time = time.time()
end_time = time.time()
last_log_time = start_time
while (end_time - start_time) < timeout:
if not thread.is_alive():
break
if end_time - last_log_time > 30:
# log every 30 seconds to avoid false detection by hangWatcher
logger.info(f"Still waiting for download, already passed {end_time - start_time} seconds...")
last_log_time = end_time
time.sleep(1)
end_time = time.time()
if thread.is_alive():
return PDCErrorCode.CommandTimeout
result = result_queue.get()
if isinstance(result, str):
logger.error(f"Unknown exception occurred during download process, details: {result}")
return PDCErrorCode.UnknownError
return result
def _pdc_download_impl(self, remote_path: str, local_path: str) -> PDCErrorCode:
"""download data from afs/bos
Args:
remote_path str: the remote file path, afs/bos, such as afs://user/a/b/xx.tar
local_path str: local file directory
Return:
"""
pre_check_status = self._pre_check()
if pre_check_status != PDCErrorCode.Success:
return pre_check_status
# check local path
if os.path.exists(local_path):
logger.info(f"local path {local_path} already exists")
return PDCErrorCode.LocalPathExist
if not remote_path.endswith(".tar"):
logger.warning(f"remote path {remote_path} should end with .tar")
return PDCErrorCode.InvalidArgument
try:
remote_dir = os.path.dirname(remote_path)
file_name = os.path.basename(remote_path)
# download remote file to local tmp path
local_tmp_file_path = f".tmp_{time.time()}_{file_name}"
status = self._download_file(remote_path, local_tmp_file_path)
if status != PDCErrorCode.Success:
logger.error(f"download remote file {file_name} failed")
return status
# download hash file to local path
hash_file_name = self._get_file_hash_name(file_name)
hash_file_path = os.path.join(remote_dir, hash_file_name)
status = self._download_file(hash_file_path, hash_file_name)
if status != PDCErrorCode.Success:
logger.error(f"download remote hash file {hash_file_path} failed")
return status
remote_hash = ""
with open(hash_file_name, "r") as f:
remote_hash = f.read().strip()
# calc hash
local_hash, status = self._calculate_hash(local_tmp_file_path)
if status != PDCErrorCode.Success:
logger.error(f"calculate hash for {local_tmp_file_path} failed")
return status
logger.info(f"remote hash: {remote_hash}, local hash: {local_hash}")
# check hash
if local_hash != remote_hash:
logger.error(f"local b3sum hash: {local_hash}, remote b3sum hash: {remote_hash}")
return PDCErrorCode.CalculateHashFail
# untar file to local_path
status = self._untar_file(local_tmp_file_path, local_path)
if status != PDCErrorCode.Success:
logger.error(f"untar file {local_tmp_file_path} failed")
return status
# clean tmp files
self._clean_tmp_files([local_tmp_file_path])
return PDCErrorCode.Success
except Exception as e:
logger.error(f"pdc upload failed: {e}")
raise e
def _pdc_download_checkpoint_impl(self, step: int) -> PDCErrorCode:
""" "download checkpoint from afs/bos
Args:
step int: the step of checkpoint
"""
pre_check_status = self._pre_check()
if pre_check_status != PDCErrorCode.Success:
return pre_check_status
conf = json.dumps(
{
"download_step": step,
}
)
# download file from remote path
download_cmd_args = [
self._pdc_agent_bin,
"-mode",
"command",
"-type",
"download_checkpoint",
"-config",
f"{conf}",
]
try:
self._pre_check()
logger.info(f"begin to download checkpoint from step {step}, config: {conf}")
res, error_code = self._exec_cmd(download_cmd_args)
if error_code == PDCErrorCode.Success:
logger.info(f"download checkpoint from step {step} successfully")
return error_code
except Exception as e:
logger.error(f"exec cmd {download_cmd_args} with error: {e}")
raise Exception(f"exec cmd {download_cmd_args} with error: {e}")
def _pre_check(self) -> PDCErrorCode:
"""check whether the environment is ready"""
if not os.path.exists(self._pdc_agent_bin):
logger.error(f"pdc tool {self._pdc_agent_bin} not found")
return PDCErrorCode.AFSToolsNotExist
if not os.path.exists(self._hash_sum_bin):
logger.error(f"hash tool {self._hash_sum_bin} not found")
return PDCErrorCode.AFSToolsNotExist
if not os.path.exists(self._train_config):
logger.error(f"train config {self._train_config} not found")
return PDCErrorCode.TrainConfigNotExist
# TODO(@zezhao): add more check
return PDCErrorCode.Success
def _exec_cmd(self, cmd_args: List[str]) -> (str, PDCErrorCode):
"""exec user command
Args:
cmd List[str]: command
"""
error_code = PDCErrorCode.Success
try:
result = subprocess.run(cmd_args, capture_output=True, text=True)
if result.returncode == 0:
logger.info(f"exec cmd {cmd_args} successfully, result: {result.stdout}; {result.stderr}")
else:
logger.error(f"exec cmd {cmd_args} failed, exit code: {result.returncode}, err: {result.stderr}")
# TODO(@zezhao): add more error code
error_code = PDCErrorCode.CommandFail
return result.stdout, error_code
except Exception as e:
logger.error(f"exec cmd {cmd_args} with error: {e}")
raise Exception(f"exec cmd {cmd_args} with error: {e}")
def _get_file_hash_name(self, file_name: str) -> str:
"""get the hash name of file
Args:
file_name str: file name
Return:
"""
if file_name.endswith(".tar"):
file_name = file_name[:-4]
return f"{file_name}.b3sumhash"
def _calculate_hash(self, file_path: str) -> (str, PDCErrorCode):
"""calc the hash of file using b3sum
Args:
file_path str: file path
Return:
"""
cmd_args = [self._hash_sum_bin, "--num-threads", "16", file_path]
try:
result, error_code = self._exec_cmd(cmd_args)
if error_code == PDCErrorCode.Success and len(result) > 0:
return result.split(" ")[0].strip(), error_code
except Exception as e:
logger.error(f"exec cmd {cmd_args} with error: {e}")
raise Exception(f"exec cmd {cmd_args} with error: {e}")
return "", PDCErrorCode.CalculateHashFail
def _tar_file(self, source_path: str, target_path: str) -> PDCErrorCode:
"""tar file with command
tar -cf target_path -C source_path .
Args:
source_path str: source file path for tar
target_path str: target file path
"""
if not os.path.exists(source_path):
logger.error(f"file {source_path} not exist")
return PDCErrorCode.LocalPathNotExist
if os.path.exists(target_path):
os.rename(target_path, f"{target_path}.old")
logger.warning(f"{target_path} already exists, backup it")
error_code = PDCErrorCode.Success
# tar file
tar_cmd_args = [self._tar_bin, "-cf", target_path, "-C", source_path, "."]
try:
res, error_code = self._exec_cmd(tar_cmd_args)
if error_code == PDCErrorCode.Success:
logger.info(f"tar {source_path} successfully")
except Exception as e:
logger.error(f"exec cmd {tar_cmd_args} failed, error: {e}")
raise Exception(f"exec cmd {tar_cmd_args} failed, error: {e}")
return error_code
def _untar_file(self, source_path: str, target_path: str) -> PDCErrorCode:
"""untar file
Args:
source_path str: source file path for untar
target_path str: target file path
"""
if not os.path.exists(source_path):
logger.error(f"{source_path} not exist")
return PDCErrorCode.LocalPathNotExist
if not os.path.exists(target_path):
# create target path if not exists
os.makedirs(target_path)
# untar file
error_code = PDCErrorCode.Success
untar_cmd_args = [self._tar_bin, "-xf", source_path, "-C", target_path]
try:
res, error_code = self._exec_cmd(untar_cmd_args)
if error_code == PDCErrorCode.Success:
logger.info(f"untar {source_path} successfully")
except Exception as e:
logger.error(f"exec cmd {untar_cmd_args} with error: {e}")
raise Exception(f"exec cmd {untar_cmd_args} with error: {e}")
return error_code
def _upload_file(self, local_file_path: str, remote_path: str) -> PDCErrorCode:
"""upload file
Args:
local_file_path str: local file path
remote_path str: remote file path
"""
if not os.path.exists(local_file_path):
logger.error(f"{local_file_path} not exist")
return PDCErrorCode.LocalPathNotExist
conf = json.dumps({"remote_path": remote_path, "local_path": local_file_path})
# upload file to remote path
upload_cmd_args = [self._pdc_agent_bin, "-mode", "command", "-type", "upload", "-config", f"{conf}"]
error_code = PDCErrorCode.Success
try:
res, error_code = self._exec_cmd(upload_cmd_args)
if error_code == PDCErrorCode.Success:
logger.info(f"upload {local_file_path} successfully")
except Exception as e:
logger.error(f"exec cmd {upload_cmd_args} with error: {e}")
raise Exception(f"exec cmd {upload_cmd_args} with error: {e}")
return error_code
def _download_file(self, remote_path: str, local_path: str) -> PDCErrorCode:
"""download file
Args:
remote_path str: remote file path
local_path str: local file path
"""
if os.path.exists(local_path):
os.rename(local_path, f"{local_path}.old")
logger.warning(f"{local_path} already exists, backup it to {local_path}.old")
conf = json.dumps({"remote_path": remote_path, "local_path": local_path})
# download file from remote path
download_cmd_args = [self._pdc_agent_bin, "-mode", "command", "-type", "download", "-config", f"{conf}"]
error_code = PDCErrorCode.Success
try:
logger.info(f"begin to download {remote_path}, config: {conf}")
res, error_code = self._exec_cmd(download_cmd_args)
if error_code == PDCErrorCode.Success:
logger.info(f"download {remote_path} successfully")
except Exception as e:
logger.error(f"exec cmd {download_cmd_args} with error: {e}")
raise Exception(f"exec cmd {download_cmd_args} with error: {e}")
return error_code
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} succeeded.")
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_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
try:
logger.info(f"begin to generate_sum path: {path}")
res, error_code = self._exec_cmd(generate_checksum_args)
if error_code == PDCErrorCode.Success:
logger.info(f"generate_sum {path} successfully")
except Exception as e:
logger.error(f"exec cmd {generate_checksum_args} with error: {e}")
return PDCErrorCode.CheckSumCommandFail
return error_code
def pdc_flash_do_check(self, path: str) -> PDCErrorCode:
"""
Args
:param localPath:
:return:
"""
if not os.path.exists(path):
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
try:
logger.info(f"begin to check_sum path: {path}")
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"[Error] [pdc_sdk] exec cmd {generate_checksum_args} with error: {e}")
self._pdc_backup_failed_directory(path)
return PDCErrorCode.CheckSumCommandFail
return error_code
def _clean_tmp_files(self, tmp_files: List[str]):
"""clean tmp files
Args:
tmp_files List[str]: list of tmp file paths
"""
if len(tmp_files) == 0:
return
# clean files
for file_path in tmp_files:
if os.path.exists(file_path):
logger.info(f"clean tmp file: {file_path}")
os.remove(file_path)
pdc_tool = PDCTools()
PDCErrorMessageMap = {
PDCErrorCode.Success: "success",
PDCErrorCode.RemotePathNotExist: "remote path not exist",
PDCErrorCode.LocalPathExist: "local path exist",
PDCErrorCode.DownloadFail: "download fail",
PDCErrorCode.AgentConfigInvalid: "agent config invalid",
PDCErrorCode.AFSToolsNotExist: "afs tools not exist",
PDCErrorCode.TrainConfigNotExist: "train config not exist",
PDCErrorCode.LocalPathNotExist: "local path not exist",
PDCErrorCode.CommandFail: "pdc agent command fail",
PDCErrorCode.CalculateHashFail: "calculate hash fail",
PDCErrorCode.InvalidArgument: "invalid argument",
PDCErrorCode.CommandTimeout: "pdc agent command timeout",
PDCErrorCode.CheckSumCommandFail: "checksum command fail",
PDCErrorCode.CopyTreeFailed: "copy directory failed",
}