Fix: Epath implements walk and glob, including tests. Fix sqlite tmp path handling. Fix S3 emulator timestamp handling

Author: voegtlel

src/megatron/energon/epathlib/epath.py

@@ -266,17 +266,15 @@ def url(self) -> str:
         return f"msc://{self.profile}{int_path_str}"
 
     def is_local(self) -> bool:
-        if self.profile == "dss":
-            # For now, a DSS path is always considered local.
-            # Note that this does not mean it exists on the local filesystem.
-            return True
-        else:
-            return self.profile == DEFAULT_PROFILE_NAME
+        # It will return a posix path if the fs is local, otherwise None
+        return self.fs.get_posix_path(self._internal_str_path) is not None
 
     def local_path(self) -> PathlibPath:
-        if not self.is_local():
+        # This resolves the path if it exists, probably ok.
+        posix_path = self.fs.get_posix_path(self._internal_str_path)
+        if posix_path is None:
             raise ValueError(f"Path {self} is not local")
-        return PathlibPath(self._internal_str_path)
+        return PathlibPath(posix_path)
 
     def is_dir(self) -> bool:
         try:
@@ -290,25 +288,46 @@ def is_file(self) -> bool:
     def mkdir(self, exist_ok: bool = True, parents: bool = False):
         pass
 
-    def glob(self, pattern) -> Generator["EPath", None, None]:
+    def walk(self) -> Generator["EPath", None, None]:
+        """Returns all files within this path (no folders)."""
+        # Prefix to be removed from found paths to remap to relative paths
+        root_prefix = self._internal_str_path.lstrip("/")
+
+        for obj in self.fs.list_recursive(self._internal_str_path):
+            rel = obj.key
+            if root_prefix:
+                if rel.startswith(root_prefix + "/"):
+                    rel = rel[len(root_prefix) + 1 :]
+                elif rel.startswith("/" + root_prefix + "/"):
+                    rel = rel[len(root_prefix) + 2 :]
+                elif rel == root_prefix or rel == "/" + root_prefix:
+                    rel = "."
+
+            path = EPath(self)
+            path.internal_path = self._resolve(self.internal_path / PurePosixPath(rel))
+            yield path
+
+    def glob(self, pattern: str) -> Generator["EPath", None, None]:
+        """Returns all files matching the pattern within this path (no folders)."""
         search_path_pattern = (self / pattern)._internal_str_path
-        # MSC S3 glob matches keys like ``bucket/key``; a leading ``/`` breaks wcmatch (pattern
+        # MSC glob matches keys like ``bucket/key``; a leading ``/`` breaks wcmatch (pattern
         # ``/b/**`` never matches ``b/parts/x``). Returned keys may repeat the bucket prefix; strip
         # it before joining with ``internal_path`` so we do not get ``/b/b/parts/...``.
-        if not self.is_local() and search_path_pattern.startswith("/"):
-            search_path_pattern = search_path_pattern.lstrip("/")
+        search_path_pattern = search_path_pattern.lstrip("/")
 
-        root_prefix = str(self.internal_path).lstrip("/")
+        # Prefix to be removed from found paths to remap to relative paths
+        root_prefix = self._internal_str_path.lstrip("/")
 
         for path in self.fs.glob(search_path_pattern):
             assert isinstance(path, str)
 
             rel = path
-            if not self.is_local() and root_prefix:
-                pfx = root_prefix + "/"
-                if rel.startswith(pfx):
-                    rel = rel[len(pfx) :]
-                elif rel == root_prefix:
+            if root_prefix:
+                if rel.startswith(root_prefix + "/"):
+                    rel = rel[len(root_prefix) + 1 :]
+                elif rel.startswith("/" + root_prefix + "/"):
+                    rel = rel[len(root_prefix) + 2 :]
+                elif rel == root_prefix or rel == "/" + root_prefix:
                     rel = "."
 
             new_path = EPath(self)

src/megatron/energon/flavors/webdataset/prepare.py

@@ -521,7 +521,7 @@ def prepare_dataset(
             tar_index_only: Only create tar-index, then exit
             media_filter: Media filter configuration
             fix_duplicates: If True, fix duplicate keys in the dataset by renaming the files in the shards.
-            index_sqlite_tmp_path: When ``parent_path`` is remote, directory used to build ``index.sqlite``
+            index_sqlite_tmp_path: When ``parent_path`` is remote, temp file path used to build ``index.sqlite``
                 locally before upload. If omitted, a new directory under ``/tmp`` is created and removed
                 after a successful run.
 
@@ -784,7 +784,17 @@ def add_media_metadata(
         progress_fn: Callable[[Iterator[Any], int], Iterator[T]] = (lambda x, y: x),
         index_sqlite_tmp_path: Optional[Path] = None,
     ) -> int:
-        """Add or refresh media metadata in an existing WebDataset index."""
+        """Add or refresh media metadata in an existing WebDataset index.
+
+        Args:
+            parent_path: WebDataset root path.
+            media_filter: Media filtering configuration.
+            workers: Number of parallel workers.
+            progress_fn: Callback for progress updates.
+            index_sqlite_tmp_path: When ``parent_path`` is remote, sqlite file path used to build
+                ``index.sqlite`` locally before upload. If omitted, a new directory under
+                ``/tmp`` is created and removed after a successful run.
+        """
 
         parent_path = EPath(parent_path)
 

src/megatron/energon/media/filesystem_prepare.py

@@ -3,7 +3,6 @@
 
 from __future__ import annotations
 
-import os
 import shutil
 import tempfile
 from functools import partial
@@ -41,6 +40,9 @@ def prepare_filesystem_dataset(
         root_path: Dataset root directory.
         media_filter: Media filtering configuration.
         progress: Whether to display a tqdm progress bar.
+        index_sqlite_tmp_path: When ``root_path`` is remote, temp file path used to build
+            ``index.sqlite`` locally before upload. If omitted, a new directory under
+            ``/tmp`` is created and removed after a successful run.
 
     Returns:
         Number of metadata entries written to the database.
@@ -137,23 +139,23 @@ def _collect_media_files(
 
         progress_bar = tqdm(total=None, unit="file", desc="Collecting media files")
 
-    if root.is_local():
-        paths = (
-            EPath(f"{path}/{file}")
-            for path, _dirs, files in os.walk(root.local_path(), followlinks=False)
-            for file in files
-        )
-    else:
-        paths = root.glob("**/*")
+    # if root.is_local() and not root.profile == "dss":
+    #     paths = (
+    #         EPath(path) / file
+    #         for path, _dirs, files in os.walk(root.local_path(), followlinks=False)
+    #         for file in files
+    #     )
+    # else:
+    #     paths = root.glob("**/*")
 
-    for file in paths:
+    for file in root.walk():
         if progress_bar is not None:
             progress_bar.update()
 
-        if not consider_all and not media_filter.should_consider_media(file.name):
+        if ("/" + MAIN_FOLDER_NAME + "/") in file.url:
             continue
 
-        if ("/" + MAIN_FOLDER_NAME + "/") in file.url:
+        if not consider_all and not media_filter.should_consider_media(file.name):
             continue
 
         files.append(file)

src/megatron/energon/tools/prepare.py

@@ -348,6 +348,11 @@ def progress_fn(els, length=None):
         def progress_fn(els, length=None):
             return els
 
+    if tmp_path is not None:
+        index_sqlite_tmp_path = tmp_path / "index.sqlite"
+    else:
+        index_sqlite_tmp_path = None
+
     found_types = BaseWebdatasetFactory.prepare_dataset(
         path,
         all_tars,
@@ -359,7 +364,7 @@ def progress_fn(els, length=None):
         workers=num_workers,
         media_filter=media_filter_config,
         fix_duplicates=fix_duplicates,
-        index_sqlite_tmp_path=tmp_path,
+        index_sqlite_tmp_path=index_sqlite_tmp_path,
     )
 
     found_types = list(found_types)

src/megatron/energon/tools/prepare_media.py

@@ -71,6 +71,11 @@ def command(
         media_metadata_by_glob, media_metadata_by_header, media_metadata_by_extension
     )
 
+    if tmp_path is not None:
+        index_sqlite_tmp_path = tmp_path / "index.sqlite"
+    else:
+        index_sqlite_tmp_path = None
+
     ds_type = get_dataset_type(path)
     if ds_type == EnergonDatasetType.WEBDATASET:
         click.echo("Preparing webdataset and computing media metadata...")
@@ -96,7 +101,7 @@ def progress_fn(els, length=None):
             media_filter=media_filter_config,
             workers=num_workers,
             progress_fn=progress_fn,
-            index_sqlite_tmp_path=tmp_path,
+            index_sqlite_tmp_path=index_sqlite_tmp_path,
         )
 
         click.echo(f"Done. Stored metadata for {count} files.")
@@ -112,6 +117,7 @@ def progress_fn(els, length=None):
         media_filter_config,
         progress=progress,
         workers=num_workers,
+        index_sqlite_tmp_path=index_sqlite_tmp_path,
     )
     click.echo(f"Done. Stored metadata for {stored} files.")
 

tests/s3_emulator/handler.py

@@ -2,7 +2,7 @@
 # SPDX-License-Identifier: BSD-3-Clause
 import urllib.parse as _up
 from datetime import datetime, timezone
-from email.utils import formatdate
+from email.utils import format_datetime
 from hashlib import md5
 from http import HTTPStatus
 from http.server import BaseHTTPRequestHandler
@@ -130,13 +130,14 @@ def _handle_write(self):
                 return
             try:
                 data = self.server.state.copy_object(bucket, key, src_bucket, src_key)
+                last_modified = self.server.state.get_object_last_modified(bucket, key)
             except FileNotFoundError:
                 self._send_error(HTTPStatus.NOT_FOUND, "NoSuchKey")
                 return
             xml = (
                 '<?xml version="1.0" encoding="UTF-8"?>'
                 "<CopyObjectResult>"
-                f"<LastModified>{_escape_xml(formatdate(usegmt=True))}</LastModified>"
+                f"<LastModified>{_escape_xml(_s3_datetime(last_modified))}</LastModified>"
                 f"<ETag>&quot;{_escape_xml(_etag(data))}&quot;</ETag>"
                 "</CopyObjectResult>"
             ).encode()
@@ -207,6 +208,7 @@ def _handle_read(self, listing: bool, only_headers: bool = False):
 
         try:
             data = self.server.state.get_object(bucket, key)
+            last_modified = self.server.state.get_object_last_modified(bucket, key)
         except FileNotFoundError:
             self._send_error(HTTPStatus.NOT_FOUND, "Not found")
             return
@@ -234,10 +236,10 @@ def _handle_read(self, listing: bool, only_headers: bool = False):
                 "Accept-Ranges": "bytes",
                 "Content-Length": str(len(slice_data)),
                 "ETag": _etag(data),
+                "Last-Modified": _http_datetime(last_modified),
             }
             if only_headers:
                 headers.setdefault("Content-Type", "application/octet-stream")
-                headers.setdefault("Last-Modified", formatdate(usegmt=True))
                 self._send_status(HTTPStatus.PARTIAL_CONTENT, extra_headers=headers)
             else:
                 self._send_bytes(
@@ -254,15 +256,19 @@ def _handle_read(self, listing: bool, only_headers: bool = False):
                         "Content-Length": str(len(data)),
                         "Accept-Ranges": "bytes",
                         "Content-Type": "application/octet-stream",
-                        "Last-Modified": formatdate(usegmt=True),
+                        "Last-Modified": _http_datetime(last_modified),
                         "ETag": _etag(data),
                     },
                 )
             else:
                 self._send_bytes(
                     data,
                     content_type="application/octet-stream",
-                    extra_headers={"Accept-Ranges": "bytes"},
+                    extra_headers={
+                        "Accept-Ranges": "bytes",
+                        "Last-Modified": _http_datetime(last_modified),
+                        "ETag": _etag(data),
+                    },
                 )
 
     def _handle_delete(self):
@@ -414,7 +420,6 @@ def _render_list_bucket_result(
         start_after = (qs.get("start-after") or [""])[0]
         exclusive_after = continuation or start_after
 
-        now = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%S.000Z")
         state = self.server.state
 
         items: list[tuple[Literal["cp", "key"], str]] = []
@@ -471,15 +476,17 @@ def _render_list_bucket_result(
             else:
                 try:
                     data = state.get_object(bucket, path)
+                    last_modified = state.get_object_last_modified(bucket, path)
                     size = len(data)
                     etag = _etag(data)
                 except Exception:  # noqa: BLE001
                     size = 0
                     etag = '""'
+                    last_modified = datetime.fromtimestamp(0, tz=timezone.utc)
                 fragments.append(
                     "<Contents>"
                     f"<Key>{_escape_xml(path)}</Key>"
-                    f"<LastModified>{now}</LastModified>"
+                    f"<LastModified>{_s3_datetime(last_modified)}</LastModified>"
                     f"<ETag>{etag}</ETag>"
                     f"<Size>{size}</Size>"
                     "</Contents>"
@@ -541,6 +548,18 @@ def _escape_xml(text: str) -> str:  # noqa: D401
     )
 
 
+def _http_datetime(value: datetime) -> str:
+    """Format an aware datetime for HTTP Last-Modified headers."""
+
+    return format_datetime(value.astimezone(timezone.utc), usegmt=True)
+
+
+def _s3_datetime(value: datetime) -> str:
+    """Format an aware datetime for S3 XML LastModified fields."""
+
+    return value.astimezone(timezone.utc).strftime("%Y-%m-%dT%H:%M:%S.000Z")
+
+
 def _etag(data: bytes) -> str:  # noqa: D401
     """Generate an ETag for binary data.