Append dataset name to download dir.

PiperOrigin-RevId: 681832795

Author: fineguy

tensorflow_datasets/core/download/download_manager.py

@@ -316,13 +316,24 @@ def downloaded_size(self) -> int:
     return sum(url_info.size for url_info in self._recorded_url_infos.values())
 
   def _get_dl_path(
-      self, resource: resource_lib.Resource, checksum: str | None = None
+      self,
+      resource: resource_lib.Resource,
+      checksum: str | None = None,
+      legacy_mode: bool = False,
   ) -> epath.Path:
-    return (
-        self._download_dir
-        / resource.relative_download_dir
-        / resource_lib.get_dl_fname(resource.url, checksum)
-    )
+    """Returns the path where the resource should be downloaded.
+
+    Args:
+      resource: The resource to download.
+      checksum: The checksum of the resource.
+      legacy_mode: If True, returns path in the legacy format without dataset
+        name in the path.
+    """
+    download_dir = self._download_dir
+    if not legacy_mode:
+      download_dir /= self._dataset_name
+    download_dir /= resource.relative_download_dir
+    return download_dir / resource_lib.get_dl_fname(resource.url, checksum)
 
   @property
   def register_checksums(self):
@@ -353,6 +364,50 @@ def _get_manually_downloaded_path(
 
     return manual_path
 
+  def _get_checksum_dl_result(
+      self, resource: resource_lib.Resource, legacy_mode: bool = False
+  ) -> downloader.DownloadResult | None:
+    """Checks if the download has been cached and checksum is known."""
+    expected_url_info = self._url_infos.get(resource.url)
+
+    if not expected_url_info:
+      return None
+
+    checksum_path = self._get_dl_path(
+        resource, expected_url_info.checksum, legacy_mode=legacy_mode
+    )
+    if not resource_lib.is_locally_cached(checksum_path):
+      return None
+
+    return downloader.DownloadResult(
+        path=checksum_path, url_info=expected_url_info
+    )
+
+  def _get_url_dl_result(
+      self, resource: resource_lib.Resource, legacy_mode: bool = False
+  ) -> downloader.DownloadResult | None:
+    """Checks if the download has been cached and checksum is unknown."""
+    url_path = self._get_dl_path(resource, legacy_mode=legacy_mode)
+    if not resource_lib.is_locally_cached(url_path):
+      return None
+
+    expected_url_info = self._url_infos.get(resource.url)
+    url_info = downloader.read_url_info(url_path)
+
+    if expected_url_info and expected_url_info != url_info:
+      # If checksums are registered but do not match, trigger a new
+      # download (e.g. previous file corrupted, checksums updated)
+      return None
+    elif self._is_checksum_registered(url=resource.url):
+      # Checksums were registered: Rename -> checksum_path
+      path = self._get_dl_path(resource, url_info.checksum)
+      resource_lib.replace_info_file(url_path, path)
+      url_path.replace(path)
+      return downloader.DownloadResult(path=path, url_info=url_info)
+    else:
+      # Checksums not registered: -> do nothing
+      return downloader.DownloadResult(path=url_path, url_info=url_info)
+
   # Synchronize and memoize decorators ensure same resource will only be
   # processed once, even if passed twice to download_manager.
   @utils.build_synchronize_decorator()
@@ -399,32 +454,20 @@ def _download_or_get_cache(
       dl_result = None
 
     # Download has been cached (checksum known)
-    elif expected_url_info and resource_lib.is_locally_cached(
-        checksum_path := self._get_dl_path(resource, expected_url_info.checksum)
-    ):
-      dl_result = downloader.DownloadResult(
-          path=checksum_path, url_info=expected_url_info
-      )
+    elif dl_result := self._get_checksum_dl_result(resource):
+      pass
+
+    # Download has been cached (checksum known, legacy mode)
+    elif dl_result := self._get_checksum_dl_result(resource, legacy_mode=True):
+      pass
 
     # Download has been cached (checksum unknown)
-    elif resource_lib.is_locally_cached(
-        url_path := self._get_dl_path(resource)
-    ):
-      url_info = downloader.read_url_info(url_path)
-
-      if expected_url_info and expected_url_info != url_info:
-        # If checksums are registered but do not match, trigger a new
-        # download (e.g. previous file corrupted, checksums updated)
-        dl_result = None
-      elif self._is_checksum_registered(url=url):
-        # Checksums were registered: Rename -> checksum_path
-        path = self._get_dl_path(resource, url_info.checksum)
-        resource_lib.replace_info_file(url_path, path)
-        url_path.replace(path)
-        dl_result = downloader.DownloadResult(path=path, url_info=url_info)
-      else:
-        # Checksums not registered: -> do nothing
-        dl_result = downloader.DownloadResult(path=url_path, url_info=url_info)
+    elif dl_result := self._get_url_dl_result(resource):
+      pass
+
+    # Download has been cached (checksum unknown, legacy mode)
+    elif dl_result := self._get_url_dl_result(resource, legacy_mode=True):
+      pass
 
     # Cache not found
     else:
@@ -504,7 +547,7 @@ def _download(
     download_tmp_dir = (
         url_path.parent / f'{url_path.name}.tmp.{uuid.uuid4().hex}'
     )
-    download_tmp_dir.mkdir()
+    download_tmp_dir.mkdir(parents=True, exist_ok=True)
     logging.info(f'Downloading {url} into {download_tmp_dir}...')
     future = self._downloader.download(
         url, download_tmp_dir, verify=self._verify_ssl

tensorflow_datasets/core/download/download_manager_test.py

@@ -33,6 +33,8 @@
 TAR = resource_lib.ExtractMethod.TAR
 NO_EXTRACT = resource_lib.ExtractMethod.NO_EXTRACT
 
+_DATASET_NAME = 'mnist'
+
 _CHECKSUMS_DIR = epath.Path('/checksums')
 _CHECKSUMS_PATH = _CHECKSUMS_DIR / 'checksums.tsv'
 
@@ -60,10 +62,10 @@ def __init__(
     )
 
     self.file_name = resource_lib.get_dl_fname(self.url, self.url_info.checksum)
-    self.file_path = _DOWNLOAD_DIR / self.file_name
+    self.file_path = _DOWNLOAD_DIR / _DATASET_NAME / self.file_name
 
     self.url_name = resource_lib.get_dl_fname(self.url)
-    self.url_path = _DOWNLOAD_DIR / self.url_name
+    self.url_path = _DOWNLOAD_DIR / _DATASET_NAME / self.url_name
 
     self.manual_path = _MANUAL_DIR / name
     extract_method = resource_lib.guess_extract_method(name)
@@ -177,7 +179,7 @@ def _get_manager(
       **kwargs,
   ):
     manager = dm.DownloadManager(
-        dataset_name='mnist',
+        dataset_name=_DATASET_NAME,
         download_dir=dl_dir,
         extract_dir=extract_dir,
         manual_dir=manual_dir,