Spaces:
Sleeping
Sleeping
File size: 2,046 Bytes
bda104d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 | r"""Utils for copying data to the temporary directory of the compute node
"""
import logging
import os
from pathlib import Path
from filelock import FileLock
log = logging.getLogger(__name__)
def copy_data(origin: Path, destination: Path):
cmd = "cp -r {origin} {destination}"
os.system(cmd.format(origin=origin, destination=destination))
def copy_to_compute_node(data_path: str | Path,
local_dir: str | Path = "/local/job"):
if not isinstance(data_path, Path):
data_path = Path(data_path)
if not isinstance(local_dir, Path):
local_dir = Path(local_dir)
# If we are not in a SLURM job, this function is a no-op
job_id = os.environ.get("SLURM_JOB_ID")
if job_id is None:
return data_path
# First, we check if the data we're looking for is already in the compute node
origin_folder = None
for folder in local_dir.iterdir():
if not os.access(folder, os.R_OK):
continue
subfolder = folder / data_path.name
if subfolder.exists():
origin_folder = subfolder
break
# Set the destination folder
dest_folder = local_dir / job_id / data_path.name
if dest_folder.exists():
log.info(f"Data found in {dest_folder}.")
return dest_folder
dest_folder.mkdir()
if origin_folder is None:
# In that case, we will copy from the login node directly (slow...)
origin_folder = data_path
else:
# Wait for the origin folder to have been fully filled to copy from the compute node (much faster!)
log.info(f"Data found in {origin_folder}. Waiting for the lock to be released...")
with FileLock(origin_folder / "lock"):
pass
# Lock the destination folder and copy the data inside it
log.info(f"Copying data from {origin_folder} to {dest_folder}...")
with FileLock(dest_folder / "lock"):
copy_data(origin_folder, dest_folder)
log.info("Done.")
return dest_folder |