Fix outdated local copy. Implement prepare media for filesystem paths on remote.

Author: voegtlel

src/megatron/energon/cache/file_store.py

@@ -161,19 +161,18 @@ def get_path(self) -> str:
     def get_media_metadata(self, key: str) -> MediaMetadataBase:
         if self._media_metadata_available is None:
             try:
-                has_metadata = self.sqlite_reader.db_has_media_metadata()
+                self._media_metadata_available = self.sqlite_reader.db_has_media_metadata()
             except sqlite3.Error as exc:  # pragma: no cover - defensive
+                self._media_metadata_available = False
                 raise RuntimeError(
                     "Failed to inspect media metadata table. Re-run `energon prepare --media-metadata-by-...`."
                 ) from exc
 
-            if not has_metadata:
-                raise RuntimeError(
-                    "Media metadata is not available for this dataset. "
-                    "Run `energon prepare --media-metadata-by-...` to generate it."
-                )
-
-            self._media_metadata_available = True
+        if not self._media_metadata_available:
+            raise RuntimeError(
+                "Media metadata is not available for this dataset. "
+                "Run `energon prepare --media-metadata-by-...` to generate it."
+            )
 
         try:
             row = self.sqlite_reader.get_media_metadata(key)

src/megatron/energon/epathlib/epath.py

@@ -336,6 +336,9 @@ def relative_to(self, other: "EPath") -> str:
 
         return str(self.internal_path.relative_to(other.internal_path))
 
+    def stat(self) -> msc.types.ObjectMetadata:
+        return self.fs.info(self._internal_str_path)
+
     @property
     def display_name(self) -> str:
         if self.profile == "dss":

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

@@ -49,6 +49,16 @@ def __init__(
                          part_name TEXT,
                          content_byte_offset INTEGER,
                          content_byte_size INTEGER)
+        if enable_media_metadata is True, it also creates the media_metadata table:
+          - media_metadata(entry_key TEXT PRIMARY KEY,
+                           metadata_type TEXT NOT NULL,
+                           metadata_json TEXT NOT NULL)
+        if enable_media_metadata is True, it also creates the media_filters table:
+          - media_filters(filter_id INTEGER PRIMARY KEY AUTOINCREMENT,
+                          strategy TEXT NOT NULL,
+                          patterns TEXT,
+                          created_at_utc TEXT DEFAULT CURRENT_TIMESTAMP,
+                          UNIQUE(strategy, patterns))
         Also creates indexes:
           - samples(sample_key)
           - samples(tar_file_id, sample_index)
@@ -309,6 +319,9 @@ class SqliteIndexReader:
                    part_name TEXT,
                    content_byte_offset INTEGER,
                    content_byte_size INTEGER)
+    - media_metadata(entry_key TEXT PRIMARY KEY,
+                     metadata_type TEXT NOT NULL,
+                     metadata_json TEXT NOT NULL)
     """
 
     sqlite_path: EPath

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

@@ -158,6 +158,8 @@ def on_start(self, aggregator_pool: AggregatorPool) -> None:
         local_sqlite = self.sqlite_path
         if self.sqlite_local_build_path is not None:
             local_sqlite = EPath(self.sqlite_local_build_path)
+            if self.sqlite_path.is_file():
+                self.sqlite_path.copy(local_sqlite)
         self.writer = SqliteIndexWriter(
             local_sqlite,
             enable_sample_tables=self.enable_sample_tables,
@@ -565,7 +567,9 @@ def prepare_dataset(
         remote_sqlite_tmp_dir: Optional[Path] = None
         if not parent_path.is_local():
             if index_sqlite_tmp_path is None:
-                remote_sqlite_tmp_dir = Path(tempfile.mkdtemp(dir="/tmp", prefix="energon-prepare-"))
+                remote_sqlite_tmp_dir = Path(
+                    tempfile.mkdtemp(dir="/tmp", prefix="energon-prepare-")
+                )
                 index_sqlite_tmp_path = remote_sqlite_tmp_dir / INDEX_SQLITE_FILENAME
                 owns_remote_sqlite_tmp = True
         else:
@@ -618,7 +622,9 @@ def prepare_dataset(
             # Fix permissions if needed
             if fix_local_permissions:
                 try:
-                    Path(str(parent_path / MAIN_FOLDER_NAME / INDEX_SQLITE_FILENAME)).chmod(file_perms)
+                    Path(str(parent_path / MAIN_FOLDER_NAME / INDEX_SQLITE_FILENAME)).chmod(
+                        file_perms
+                    )
                 except OSError:
                     pass
 
@@ -703,7 +709,9 @@ def prepare_dataset(
                 for split_part, split_ratio in split_parts_ratio:
                     split_total += split_ratio
                     split_end = int(len(shards) * split_total)
-                    split_shards[split_part] = [shard.name for shard in shards[split_offset:split_end]]
+                    split_shards[split_part] = [
+                        shard.name for shard in shards[split_offset:split_end]
+                    ]
                     split_offset = split_end
             else:
                 assert split_parts_patterns is not None, (
@@ -749,7 +757,6 @@ def prepare_dataset(
             if owns_remote_sqlite_tmp and remote_sqlite_tmp_dir is not None:
                 shutil.rmtree(remote_sqlite_tmp_dir, ignore_errors=True)
 
-
     @classmethod
     def add_media_metadata(
         cls,
@@ -781,12 +788,16 @@ def add_media_metadata(
         remote_sqlite_tmp_dir: Optional[Path] = None
         if not parent_path.is_local():
             if index_sqlite_tmp_path is None:
-                remote_sqlite_tmp_dir = Path(tempfile.mkdtemp(dir="/tmp", prefix="energon-prepare-media-"))
+                remote_sqlite_tmp_dir = Path(
+                    tempfile.mkdtemp(dir="/tmp", prefix="energon-prepare-media-")
+                )
                 index_sqlite_tmp_path = remote_sqlite_tmp_dir / INDEX_SQLITE_FILENAME
                 owns_remote_sqlite_tmp = True
         else:
             index_sqlite_tmp_path = None
 
+        sqlite_path = parent_path / MAIN_FOLDER_NAME / INDEX_SQLITE_FILENAME
+
         try:
             aggregator = SqliteIndexWriterAggregator(
                 parent_path / MAIN_FOLDER_NAME / INDEX_SQLITE_FILENAME,
@@ -818,6 +829,18 @@ def add_media_metadata(
 
             pool.process()
 
+            if sqlite_path.is_local():
+                try:
+                    meta_dir = (parent_path / MAIN_FOLDER_NAME).local_path()
+                    # Copy group permissions from the parent dir
+                    meta_dir.chmod((parent_path.local_path().stat().st_mode | 0o700))
+                    # Just read/write, no execute
+                    sqlite_path.local_path().chmod(
+                        (parent_path.local_path().stat().st_mode | 0o600) & 0o666
+                    )
+                except OSError:
+                    pass
+
             return aggregator.media_metadata_written
         finally:
             if owns_remote_sqlite_tmp and remote_sqlite_tmp_dir is not None:

src/megatron/energon/local_copy.py

@@ -3,18 +3,20 @@
 
 
 import hashlib
+import os
 from pathlib import Path
 
 from filelock import FileLock
 
 from megatron.energon.epathlib import EPath
 
-LOCAL_COPY_TMP_DIR = Path("/tmp/energon_local_copy")
+LOCAL_COPY_TMP_DIR = Path(os.environ.get("ENERGON_LOCAL_COPY_TMP_DIR", "/tmp/energon_local_copy"))
 
 
 def ensure_local_copy(path: EPath) -> EPath:
     """If the path is not local, copy it to a temporary directory and return the
-    path to the temporary directory.
+    path to the temporary directory. Assuming that the local file is never modified
+    after it is copied. Will re-sync if the remote file is newer.
 
     Args:
         path: The path to the file to copy.
@@ -33,10 +35,13 @@ def ensure_local_copy(path: EPath) -> EPath:
     final_path = LOCAL_COPY_TMP_DIR / f"{digest}.bin"
     lock_path = final_path.with_suffix(".lock")
     tmp_path = final_path.with_suffix(".part")
+    mod_time = path.stat().last_modified.timestamp()
 
     # Block until lock is free
     with FileLock(lock_path, timeout=60 * 5):
-        if final_path.exists():  # someone else already produced it
+        # someone else already produced it
+        if final_path.exists() and final_path.stat().st_mtime >= mod_time:
+            # The local file is already newer than the remote file
             return EPath(final_path)
 
         # We are the downloader

src/megatron/energon/media/extractor.py

@@ -15,12 +15,13 @@
 from PIL import Image, UnidentifiedImageError
 
 from megatron.energon.av import AVDecoder
+from megatron.energon.epathlib import EPath
 from megatron.energon.media.metadata import ImageMetadata, MediaMetadataBase, MediaMetadataType
 
 logger = logging.getLogger(__name__)
 
 
-SourceData = Union[bytes, Path, BinaryIO]
+SourceData = Union[bytes, EPath, BinaryIO]
 
 
 class MediaFilterStrategy(str, Enum):
@@ -210,9 +211,13 @@ def _build_metadata(
 
 
 def _build_image_metadata(source: SourceData) -> ImageMetadata | None:
+    should_close = False
     try:
         if isinstance(source, (bytes, bytearray)):
             source = io.BytesIO(source)
+        elif isinstance(source, EPath):
+            source = source.open("rb")
+            should_close = True
 
         with Image.open(source) as image:
             image.load()
@@ -225,6 +230,9 @@ def _build_image_metadata(source: SourceData) -> ImageMetadata | None:
     except UnidentifiedImageError:
         logger.debug("Failed to parse image metadata", exc_info=True)
         return None
+    finally:
+        if should_close:
+            source.close()
 
 
 def _build_av_metadata(source: SourceData) -> MediaMetadataBase | None: