Create check_dask_cwd function (NVIDIA#484)

* Create `check_dask_cwd` function Signed-off-by: Sarah Yurick <[email protected]> * add Praateek's suggestions Signed-off-by: Sarah Yurick <[email protected]> --------- Signed-off-by: Sarah Yurick <[email protected]>
sarahyurick · Jan 21, 2025 · 57f0e3c · 57f0e3c
1 parent 2a39616
commit 57f0e3c
Showing 1 changed file with 22 additions and 0 deletions.
diff --git a/nemo_curator/utils/distributed_utils.py b/nemo_curator/utils/distributed_utils.py
@@ -16,6 +16,7 @@
 import ast
 import os
 import shutil
+import subprocess
 
 import dask
 
@@ -565,6 +566,9 @@ def read_data(
     """
     if isinstance(input_files, str):
         input_files = [input_files]
+
+    check_dask_cwd(input_files)
+
     if file_type == "pickle":
         df = read_pandas_pickle(
             input_files[0], add_filename=add_filename, columns=columns, **kwargs
@@ -1013,6 +1017,24 @@ def get_current_client():
         return None
 
 
+def check_dask_cwd(file_list: List[str]):
+    if any(not os.path.isabs(file_path) for file_path in file_list):
+        dask_cwd_list = list(get_current_client().run(os.getcwd).values())
+        if len(set(dask_cwd_list)) <= 1:
+            dask_cwd = dask_cwd_list[0]
+            os_pwd = subprocess.check_output("pwd", shell=True, text=True).strip()
+            if dask_cwd != os_pwd:
+                raise RuntimeError(
+                    "Mismatch between Dask client and worker working directories. "
+                    "Use absolute file paths to ensure the correct files are read as intended."
+                )
+        else:
+            raise RuntimeError(
+                "Mismatch between at least 2 Dask workers' working directories. "
+                "Use absolute file paths to ensure the correct files are read as intended."
+            )
+
+
 def performance_report_if(
     path: Optional[str] = None, report_name: str = "dask-profile.html"
 ):