Skip to content

fix(readers): fail fast on unresolved input paths #2352

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
mohammadaaftabv wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

fix(readers): fail fast on unresolved input paths #2352

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
272 changes: 252 additions & 20 deletions nemo_curator/stages/file_partitioning.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,9 +12,12 @@
# See the License for the specific language governing permissions and
# limitations under the License.

from dataclasses import dataclass
import os
from dataclasses import dataclass, field
from typing import Any

from fsspec import available_protocols
from fsspec.core import split_protocol, url_to_fs
from loguru import logger

from nemo_curator.stages.base import ProcessingStage
Expand All @@ -37,8 +40,13 @@ class FilePartitioningStage(ProcessingStage[EmptyTask, FileGroupTask]):

Parameters
----------
file_paths: str | list[str]
Path to the input files.
file_paths: str | os.PathLike[str] | list[str | os.PathLike[str]]
Path, glob, or list of paths/globs for the input files. Local paths are
first resolved relative to the working directory where the stage is
constructed. If that candidate contains no supported files, discovery
retries the original worker-relative path. When the stage itself has a
runtime ``working_dir``, relative paths resolve there directly. Remote
URI components are preserved.
files_per_partition: int | None = None
Number of files per partition.
If both files_per_partition and blocksize are not provided,
Expand All @@ -54,15 +62,27 @@ class FilePartitioningStage(ProcessingStage[EmptyTask, FileGroupTask]):
Storage options to pass to the file system.
limit: int | None = None
Maximum number of partitions to create.
allow_empty: bool = False
If True, return no tasks when no supported files are discovered. By
default an empty source raises ``FileNotFoundError``.

Raises
------
FileNotFoundError
If no files with the configured extensions are discovered and
``allow_empty`` is False.
"""

file_paths: str | list[str]
file_paths: str | os.PathLike[str] | list[str | os.PathLike[str]]
files_per_partition: int | None = None
blocksize: int | str | None = None
file_extensions: list[str] | None = None
storage_options: dict[str, Any] | None = None
limit: int | None = None
name: str = "file_partitioning"
allow_empty: bool = False

_construction_cwd: str = field(init=False, repr=False)

def __post_init__(self):
"""Initialize default values."""
Expand All @@ -74,6 +94,9 @@ def __post_init__(self):
if self.storage_options is None:
self.storage_options = {}

self.file_paths = self._coerce_file_paths(self.file_paths)
self._construction_cwd = os.getcwd()

# self.blocksize is the value set by the user
# self._blocksize is the value used internally
if self.blocksize is not None:
Expand Down Expand Up @@ -106,14 +129,13 @@ def process(self, _: EmptyTask) -> list[FileGroupTask]:
and outputs multiple FileGroupTasks for parallel processing.
"""
sort_by_size = self.blocksize is not None
files_with_sizes = self._get_file_list_with_sizes(sort_by_size)
files_with_sizes, attempted_paths = self._discover_file_list_with_sizes(sort_by_size)
# Extract list[str] from list[tuple[str, int]]
files = [file[0] for file in files_with_sizes]

logger.info(f"Found {len(files)} files")
if len(files) == 0:
logger.warning(f"No files found under {self.file_paths}")
return []
return self._handle_empty_source(attempted_paths)

# Partition files
if self.files_per_partition:
Expand Down Expand Up @@ -182,34 +204,244 @@ def process(self, _: EmptyTask) -> list[FileGroupTask]:
logger.info(f"Created {len(tasks)} file groups from {len(files)} files")
return tasks

def _get_file_list_with_sizes(self, sort_by_size: bool = True) -> list[tuple[str, int]]:
def _handle_empty_source(self, attempted_paths: list[str]) -> list[FileGroupTask]:
if self.allow_empty:
logger.warning(f"No files found after trying {attempted_paths}")
return []
expected_extensions = ", ".join(self.file_extensions)
msg = (
f"No supported input files were found for configured path(s) {self.file_paths!r}, "
f"after trying these discovery path(s) in order: {attempted_paths!r}. "
f"Expected file extensions: {expected_extensions}. "
"Check that the path exists, any glob pattern matches files, and the execution environment "
"can access the input."
)
raise FileNotFoundError(msg)

@staticmethod
def _coerce_path(path: str | os.PathLike[str]) -> str:
if not isinstance(path, (str, os.PathLike)):
msg = f"Invalid file path: {path!r}, must be a string or path-like object"
raise TypeError(msg)
coerced_path = os.fspath(path)
if not isinstance(coerced_path, str):
msg = f"Invalid file path: {path!r}, must resolve to a string"
raise TypeError(msg)
return coerced_path

@classmethod
def _coerce_file_paths(
cls,
file_paths: str | os.PathLike[str] | list[str | os.PathLike[str]],
) -> str | list[str]:
if isinstance(file_paths, (str, os.PathLike)):
return cls._coerce_path(file_paths)
if isinstance(file_paths, list):
return [cls._coerce_path(path) for path in file_paths]
msg = f"Invalid file paths: {file_paths!r}, must be a string, path-like object, or list of them"
raise TypeError(msg)

@staticmethod
def _normalize_local_path(path: str, *, base_dir: str, preserve_relative: bool) -> str:
expanded_path = os.path.expanduser(path)
if os.path.isabs(expanded_path):
return expanded_path
if preserve_relative:
return path
return os.path.abspath(os.path.join(base_dir, expanded_path))

@classmethod
def _normalize_input_path(cls, path: str, *, base_dir: str, preserve_relative: bool) -> str:
"""Normalize local parts of an fsspec path while retaining its protocol chain."""
components = path.split("::")
is_chain = len(components) > 1
known_protocols = set(available_protocols()) if is_chain else set()
normalized_components = []

for component in components:
protocol, protocol_path = split_protocol(component)
if protocol in {"file", "local"}:
normalized_path = cls._normalize_local_path(
protocol_path,
base_dir=base_dir,
preserve_relative=preserve_relative,
)
normalized_components.append(f"{protocol}://{normalized_path}")
elif protocol is not None:
normalized_components.append(component)
elif component.startswith(("file:", "local:")):
protocol, protocol_path = component.split(":", 1)
normalized_path = cls._normalize_local_path(
protocol_path,
base_dir=base_dir,
preserve_relative=preserve_relative,
)
normalized_components.append(f"{protocol}://{normalized_path}")
elif is_chain and component in known_protocols:
normalized_components.append(component)
else:
normalized_components.append(
cls._normalize_local_path(
component,
base_dir=base_dir,
preserve_relative=preserve_relative,
)
)

return "::".join(normalized_components)

def _get_file_paths_for_discovery(self) -> str | list[str]:
if isinstance(self.file_paths, str):
return self._get_discovery_path_candidates(self.file_paths)[0]
return [self._get_discovery_path_candidates(path)[0] for path in self.file_paths]

def _get_discovery_path_candidates(self, path: str) -> list[str]:
"""Return ordered, unique discovery candidates for one configured path."""
runtime_working_dir = self.runtime_env.get("working_dir") if self.runtime_env else None
if runtime_working_dir is not None:
return [
self._normalize_input_path(
path,
base_dir=self._construction_cwd,
preserve_relative=True,
)
]

anchored_path = self._normalize_input_path(
path,
base_dir=self._construction_cwd,
preserve_relative=False,
)
worker_relative_path = self._normalize_input_path(
path,
base_dir=self._construction_cwd,
preserve_relative=True,
)
return list(dict.fromkeys((anchored_path, worker_relative_path)))

@staticmethod
def _replace_chained_path_pattern(path: str, discovered_path: str) -> str:
"""Replace the outer filesystem's path pattern while retaining its chain targets."""
components = path.split("::")
known_protocols = set(available_protocols())

for index, component in enumerate(components):
protocol, protocol_path = split_protocol(component)
if protocol is not None:
if protocol_path:
components[index] = f"{protocol}://{discovered_path}"
return "::".join(components)
continue
if component.startswith(("file:", "local:")):
protocol, protocol_path = component.split(":", 1)
if protocol_path:
components[index] = f"{protocol}://{discovered_path}"
return "::".join(components)
continue
if component in known_protocols:
continue
components[index] = discovered_path
return "::".join(components)

msg = f"Could not locate a replaceable path component in chained URL {path!r}"
raise ValueError(msg)

def _get_file_records_for_path(
self,
path: str,
*,
recurse_subdirectories: bool,
sort_by_size: bool,
) -> list[tuple[str, int]]:
if "::" not in path:
return get_all_file_paths_and_size_under(
path,
recurse_subdirectories=recurse_subdirectories,
keep_extensions=self.file_extensions,
storage_options=self.storage_options,
sort_by_size=sort_by_size,
)

fs, fs_path = url_to_fs(path, **self.storage_options)
records = get_all_file_paths_and_size_under(
fs_path,
recurse_subdirectories=recurse_subdirectories,
keep_extensions=self.file_extensions,
fs=fs,
sort_by_size=sort_by_size,
)
return [(self._replace_chained_path_pattern(path, record_path), size) for record_path, size in records]

def _discover_file_list_with_sizes(self, sort_by_size: bool) -> tuple[list[tuple[str, int]], list[str]]:
"""Discover each configured path, falling back to worker-relative paths independently."""
configured_paths = [self.file_paths] if isinstance(self.file_paths, str) else self.file_paths
recurse_subdirectories = isinstance(self.file_paths, str)
attempted_paths: list[str] = []
records_by_path: dict[str, int] = {}
candidate_records_by_path: dict[str, list[tuple[str, int]]] = {}

for configured_path in configured_paths:
candidates = self._get_discovery_path_candidates(configured_path)
for candidate_index, candidate in enumerate(candidates):
if candidate in candidate_records_by_path:
candidate_records = candidate_records_by_path[candidate]
else:
attempted_paths.append(candidate)
try:
candidate_records = self._get_file_records_for_path(
candidate,
recurse_subdirectories=recurse_subdirectories,
sort_by_size=sort_by_size,
)
except FileNotFoundError:
candidate_records = []
candidate_records_by_path[candidate] = candidate_records

if candidate_records:
for record_path, size in candidate_records:
records_by_path.setdefault(record_path, size)
break

if candidate_index + 1 < len(candidates):
logger.debug(
f"No supported files found for {candidate!r}; "
f"retrying worker-relative candidate {candidates[candidate_index + 1]!r}"
)

records = list(records_by_path.items())
return sorted(records, key=lambda x: x[1] if sort_by_size else x[0]), attempted_paths

def _get_file_list_with_sizes(
self,
sort_by_size: bool = True,
file_paths: str | list[str] | None = None,
) -> list[tuple[str, int]]:
"""
Get the list of files to process.
"""
logger.debug(f"Getting file list with sizes for {self.file_paths}")
if isinstance(self.file_paths, str):
if file_paths is None:
records, _ = self._discover_file_list_with_sizes(sort_by_size)
return records
logger.debug(f"Getting file list with sizes for {file_paths}")
if isinstance(file_paths, str):
# Directory: list contents (recursively) and filter extensions
output_ls = get_all_file_paths_and_size_under(
self.file_paths,
output_ls = self._get_file_records_for_path(
file_paths,
recurse_subdirectories=True,
keep_extensions=self.file_extensions,
storage_options=self.storage_options,
sort_by_size=sort_by_size,
)
elif isinstance(self.file_paths, list):
elif isinstance(file_paths, list):
output_ls = []
for path in self.file_paths:
for path in file_paths:
output_ls.extend(
get_all_file_paths_and_size_under(
self._get_file_records_for_path(
path,
recurse_subdirectories=False,
keep_extensions=self.file_extensions,
storage_options=self.storage_options,
sort_by_size=sort_by_size,
)
)
else:
msg = f"Invalid file paths: {self.file_paths}, must be a string or list of strings"
msg = f"Invalid file paths: {file_paths}, must be a string or list of strings"
raise TypeError(msg)
return sorted(output_ls, key=lambda x: x[1] if sort_by_size else x[0])

Expand Down
Loading
Loading