| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import fnmatch |
| 5 | from typing import TYPE_CHECKING |
| 6 | |
| 7 | from vllm.utils.import_utils import PlaceholderModule |
| 8 | |
| 9 | if TYPE_CHECKING: |
| 10 | from botocore.client import BaseClient |
| 11 | |
| 12 | try: |
| 13 | import boto3 |
| 14 | except ImportError: |
| 15 | boto3 = PlaceholderModule("boto3") # type: ignore[assignment] |
| 16 | |
| 17 | |
| 18 | def _filter_allow(paths: list[str], patterns: list[str]) -> list[str]: |
| 19 | return [ |
| 20 | path |
| 21 | for path in paths |
| 22 | if any(fnmatch.fnmatch(path, pattern) for pattern in patterns) |
| 23 | ] |
| 24 | |
| 25 | |
| 26 | def _filter_ignore(paths: list[str], patterns: list[str]) -> list[str]: |
| 27 | return [ |
| 28 | path |
| 29 | for path in paths |
| 30 | if not any(fnmatch.fnmatch(path, pattern) for pattern in patterns) |
| 31 | ] |
| 32 | |
| 33 | |
| 34 | def glob( |
| 35 | s3: "BaseClient | None" = None, |
| 36 | path: str = "", |
| 37 | allow_pattern: list[str] | None = None, |
| 38 | ) -> list[str]: |
| 39 | """List full file names from S3 path and filter by allow pattern. |
| 40 | |
| 41 | Args: |
| 42 | s3: S3 client to use. |
| 43 | path: The S3 path to list from. |
| 44 | allow_pattern: A list of patterns of which files to pull. |
| 45 | |
| 46 | Returns: |
| 47 | list[str]: List of full S3 paths allowed by the pattern |
| 48 | |
| 49 | """ |
| 50 | if s3 is None: |
| 51 | s3 = boto3.client("s3") |
| 52 | if not path.endswith("/"): |
| 53 | path = path + "/" |
| 54 | bucket_name, _, paths = list_files(s3, path=path, allow_pattern=allow_pattern) |
| 55 | return [f"s3://{bucket_name}/{path}" for path in paths] |
| 56 | |
| 57 | |
| 58 | def list_files( |
| 59 | s3: "BaseClient", |
| 60 | path: str, |
| 61 | allow_pattern: list[str] | None = None, |
| 62 | ignore_pattern: list[str] | None = None, |
| 63 | ) -> tuple[str, str, list[str]]: |
| 64 | """List files from S3 path and filter by pattern. |
| 65 | |
| 66 | Args: |
| 67 | s3: S3 client to use. |
| 68 | path: The S3 path to list from. |
| 69 | allow_pattern: A list of patterns of which files to pull. |
| 70 | ignore_pattern: A list of patterns of which files not to pull. |
| 71 | |
| 72 | Returns: |
| 73 | tuple[str, str, list[str]]: A tuple where: |
| 74 | - The first element is the bucket name |
| 75 | - The second element is string represent the bucket |
| 76 | and the prefix as a dir like string |
| 77 | - The third element is a list of files allowed or |
| 78 | disallowed by pattern |
| 79 | |
| 80 | """ |
| 81 | parts = path.removeprefix("s3://").split("/") |
| 82 | prefix = "/".join(parts[1:]) |
| 83 | bucket_name = parts[0] |
| 84 | |
| 85 | objects = s3.list_objects_v2(Bucket=bucket_name, Prefix=prefix) |
| 86 | paths = [obj["Key"] for obj in objects.get("Contents", [])] |
| 87 | |
| 88 | paths = _filter_ignore(paths, ["*/"]) |
| 89 | if allow_pattern is not None: |
| 90 | paths = _filter_allow(paths, allow_pattern) |
| 91 | |
| 92 | if ignore_pattern is not None: |
| 93 | paths = _filter_ignore(paths, ignore_pattern) |
| 94 | |
| 95 | return bucket_name, prefix, paths |