Delete a shard file that fails its checksum so the retry downloads it again

CHANGELOG.md

@@ -10,6 +10,7 @@ streaming-community is a fork of MosaicML Streaming 0.13.0. Each entry names the
 - Raise the error of a failed download thread from iteration, with the cause attached, instead of ending the epoch early or waiting forever for the thread. In 0.14.0, a `cache_limit` too tight for a mix of local-only and remote streams hung or returned an empty epoch this way; that case now fails at construction with the limit to set. The ready thread also downloads a shard again if it was evicted while being waited for, and stops at interpreter exit, so a stalled epoch no longer blocks shutdown (#49).
 - Close each shared memory block that the prefix lookup opens to read another dataset's local directories. Every probe stayed open until the process exited, two file descriptors each, so they piled up in a process that creates many datasets (#54).
 - Return a shard to its previous state when its download or decompression raises. It stayed `PREPARING`, so every worker waiting for it waited forever and no later call downloaded it again; now the next call retries (#54).
+- Delete a downloaded shard file that fails its checksum. It stayed in the cache, so a retry, or the next read of that shard, used the corrupt file without checking it again; now the next attempt downloads it again. A local directory that holds the dataset itself is left alone (#56).
 
 ### Docs
 

streaming/base/stream.py

@@ -337,6 +337,10 @@ def _decompress_shard_part(self, zip_info: FileInfo, zip_filename: str, raw_file
                     f'validation does not match with those provided during dataset ' +
                     f'creation `{sorted(zip_info.hashes.keys())}`. Provide one of those.')
             if get_hash(self.validate_hash, data) != zip_info.hashes[self.validate_hash]:
+                # Delete a corrupt cached copy so that the next attempt downloads it again,
+                # unless the local directory holds the dataset itself.
+                if self.remote not in {None, self.local}:
+                    os.remove(zip_filename)
                 raise ValueError(f'Checksum failure: {zip_filename}')
 
         # Decompress and save that.
@@ -408,6 +412,9 @@ def _prepare_shard_part(self,
                             f'creation `{sorted(raw_info.hashes.keys())}`. Provide one of those.')
                     data = open(raw_filename, 'rb').read()
                     if get_hash(self.validate_hash, data) != raw_info.hashes[self.validate_hash]:
+                        # Delete the corrupt download so that the next attempt fetches it again.
+                        if self.remote not in {None, self.local}:
+                            os.remove(raw_filename)
                         raise ValueError(f'Checksum failure: {raw_filename}')
         return delta
 

tests/test_eviction.py

@@ -392,8 +392,11 @@ def streams() -> list[Stream]:
     assert set(os.listdir(local_only)) == set(os.listdir(remote))
 
 
-def _write_small_shards(remote: str, compression: Any = None) -> None:
-    with MDSWriter(out=remote, columns={'data': 'bytes'}, compression=compression,
+def _write_small_shards(remote: str, compression: Any = None, hashes: Any = None) -> None:
+    with MDSWriter(out=remote,
+                   columns={'data': 'bytes'},
+                   compression=compression,
+                   hashes=hashes,
                    size_limit=500) as out:
         for _ in range(200):
             out.write({'data': b'\0'})
@@ -441,3 +444,35 @@ def test_failed_decompression_keeps_shard_local(local_remote_dir: tuple[str, str
     assert dataset._shard_states[0] == _ShardState.LOCAL
     dataset.prepare_shard(0, blocking=False)
     assert os.path.isfile(raw_filename)
+
+
+@pytest.mark.parametrize('compression', [None, 'zstd'])
+def test_checksum_failure_downloads_shard_again(local_remote_dir: tuple[str, str],
+                                                compression: Any):
+    """A downloaded file that fails its checksum is deleted, so the retry downloads it again."""
+    local, remote = local_remote_dir
+    _write_small_shards(remote, compression, hashes=['sha1'])
+    dataset = StreamingDataset(remote=remote, local=local, batch_size=1, validate_hash='sha1')
+    raw_info, zip_info = dataset.shards[0].file_pairs[0]
+    basename = (zip_info or raw_info).basename
+    download = Stream._download_file
+    corrupted = []
+
+    def corrupt_first(self: Stream, from_basename: str, to_basename: Any = None) -> str:
+        filename = download(self, from_basename, to_basename)
+        if from_basename == basename and not corrupted:
+            corrupted.append(filename)
+            with open(filename, 'r+b') as f:
+                f.seek(-1, os.SEEK_END)
+                last = f.read(1)[0]
+                f.seek(-1, os.SEEK_END)
+                f.write(bytes([last ^ 0xff]))
+        return filename
+
+    with patch.object(Stream, '_download_file', corrupt_first):
+        with pytest.raises(ValueError, match='Checksum failure'):
+            dataset.prepare_shard(0)
+        assert not os.path.exists(corrupted[0])
+        dataset.prepare_shard(0)
+    assert dataset._shard_states[0] == _ShardState.LOCAL
+    assert dataset[0]['data'] == b'\0'