From dc9336d76413d51b39511bb21a5b501c0005bd9d Mon Sep 17 00:00:00 2001
From: kramstrom <kasper@modal.com>
Date: Fri, 12 Jul 2024 16:07:35 -0400
Subject: [PATCH] more progress bars!

---
 modal/_output.py                 | 4 ++--
 modal/cli/_download.py           | 3 ++-
 modal/cli/network_file_system.py | 9 ++++-----
 modal/network_file_system.py     | 3 ++-
 4 files changed, 10 insertions(+), 9 deletions(-)

diff --git a/modal/_output.py b/modal/_output.py
index 93d954dc1..bf7794669 100644
--- a/modal/_output.py
+++ b/modal/_output.py
@@ -471,10 +471,10 @@ def progress(
             if complete:
                 self._complete_sub_task(task_id)
                 return task_id
-            if advance:
+            if advance is not None:
                 self._advance_sub_task(task_id, advance)
                 return task_id
-        if name and size:
+        if name is not None and size is not None:
             return self._add_sub_task(name, size)
         raise NotImplementedError(
             "Unknown action to take with args: "
diff --git a/modal/cli/_download.py b/modal/cli/_download.py
index beaac8981..c8a4d5d6e 100644
--- a/modal/cli/_download.py
+++ b/modal/cli/_download.py
@@ -9,6 +9,7 @@
 from click import UsageError
 
 from modal._utils.async_utils import TaskContext
+from modal.config import logger
 from modal.network_file_system import _NetworkFileSystem
 from modal.volume import FileEntry, FileEntryType, _Volume
 
@@ -81,7 +82,7 @@ async def consumer():
                             async for chunk in volume.read_file(entry.path):
                                 b += fp.write(chunk)
                                 progress_cb(task_id=progress_task_id, advance=len(chunk))
-                        print(f"Wrote {b} bytes to {output_path}", file=sys.stderr)
+                        logger.debug(f"Wrote {b} bytes to {output_path}", file=sys.stderr)
                         progress_cb(task_id=progress_task_id, complete=True)
                     elif entry.type == FileEntryType.DIRECTORY:
                         output_path.mkdir(parents=True, exist_ok=True)
diff --git a/modal/cli/network_file_system.py b/modal/cli/network_file_system.py
index f9d2f3d72..263584b0b 100644
--- a/modal/cli/network_file_system.py
+++ b/modal/cli/network_file_system.py
@@ -8,14 +8,13 @@
 from click import UsageError
 from grpclib import GRPCError, Status
 from rich.console import Console
-from rich.live import Live
 from rich.syntax import Syntax
 from rich.table import Table
 from typer import Typer
 
 import modal
 from modal._location import display_location
-from modal._output import ProgressHandler, step_completed, step_progress
+from modal._output import ProgressHandler, step_completed
 from modal._utils.async_utils import synchronizer
 from modal._utils.grpc_utils import retry_transient_errors
 from modal.cli._download import _volume_download
@@ -143,9 +142,9 @@ async def put(
     console = Console()
 
     if Path(local_path).is_dir():
-        spinner = step_progress(f"Uploading directory '{local_path}' to '{remote_path}'...")
-        with Live(spinner, console=console):
-            await volume.add_local_dir(local_path, remote_path)
+        progress_handler = ProgressHandler(type="upload", console=console)
+        with progress_handler.live:
+            await volume.add_local_dir(local_path, remote_path, progress_cb=progress_handler.progress)
         console.print(step_completed(f"Uploaded directory '{local_path}' to '{remote_path}'"))
 
     elif "*" in local_path:
diff --git a/modal/network_file_system.py b/modal/network_file_system.py
index 4a0e753f1..0361407a6 100644
--- a/modal/network_file_system.py
+++ b/modal/network_file_system.py
@@ -327,6 +327,7 @@ async def add_local_dir(
         self,
         local_path: Union[Path, str],
         remote_path: Optional[Union[str, PurePosixPath, None]] = None,
+        progress_cb: Callable = None,
     ):
         _local_path = Path(local_path)
         if remote_path is None:
@@ -346,7 +347,7 @@ def gen_transfers():
         transfer_paths = aiostream.stream.iterate(gen_transfers())
         await aiostream.stream.map(
             transfer_paths,
-            aiostream.async_(lambda paths: self.add_local_file(paths[0], paths[1])),
+            aiostream.async_(lambda paths: self.add_local_file(paths[0], paths[1], progress_cb)),
             task_limit=20,
         )