more progress bars!

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: "

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)

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:

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,
         )