Fix: Hide debug statements behind _DEBUG (#730)

* Fix: Hide debug statements behind _DEBUG

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: robTheBuildr

src/litdata/streaming/downloader.py

@@ -26,6 +26,7 @@
 
 from litdata.constants import (
     _AZURE_STORAGE_AVAILABLE,
+    _DEBUG,
     _GOOGLE_STORAGE_AVAILABLE,
     _HF_HUB_AVAILABLE,
     _INDEX_FILENAME,
@@ -258,7 +259,8 @@ def download_file(self, remote_filepath: str, local_filepath: str) -> None:
                     ExtraArgs=extra_args,
                     Config=TransferConfig(use_threads=False),
                 )
-                print("DOWNLOAD TIME", time() - t0)
+                if _DEBUG:
+                    print("DOWNLOAD TIME", time() - t0)
 
     def download_bytes(self, remote_filepath: str, offset: int, length: int, local_chunkpath: str) -> bytes:
         obj = parse.urlparse(remote_filepath)

src/litdata/streaming/item_loader.py

@@ -222,7 +222,7 @@ def load_item_from_chunk(
                 if (time() - start_time) > _MAX_WAIT_TIME:
                     raise FileNotFoundError(f"The {chunk_filepath} hasn't been found.")
 
-            if time() - start_time > 5:
+            if _DEBUG and time() - start_time > 5:
                 print("WAIT TIME", time() - start_time)
 
             self._chunk_filepath = chunk_filepath
@@ -340,7 +340,8 @@ def delete(self, chunk_index: int, chunk_filepath: str) -> None:
             )
         )
         if os.path.exists(chunk_filepath):
-            print(f"delete_chunk_{chunk_index}")
+            if _DEBUG:
+                print(f"delete_chunk_{chunk_index}")
             os.remove(chunk_filepath)
         logger.debug(
             _get_log_msg(

src/litdata/streaming/reader.py

@@ -83,7 +83,7 @@ def __init__(
         num_bytes_per_nodes = self._config.num_bytes // self._distributed_env.num_nodes
         self._delete_chunks_when_processed = num_bytes_per_nodes > max_cache_size if max_cache_size else False
 
-        if distributed_env.global_rank == 0 and self._worker_env.rank == 0:
+        if _DEBUG and distributed_env.global_rank == 0 and self._worker_env.rank == 0:
             print(f"Delete chunks when used: {self._delete_chunks_when_processed}")
 
         self._has_exited = False