adding progress bar for asyncio operations

examples/basics/hello_v2_with_progress_bar.py

@@ -0,0 +1,35 @@
+from tqdm.asyncio import tqdm
+
+import flyte
+
+env = flyte.TaskEnvironment(
+    name="hello_v2_with_progress_bar",
+)
+
+
+@env.task()
+async def hello_worker(id: int) -> str:
+    ctx = flyte.ctx()
+    assert ctx is not None
+    return f"hello, my id is: {id} and I am being run by Action: {ctx.action}"
+
+
+@env.task()
+async def hello_driver(ids: list[int] = [1, 2, 3]) -> list[str]:
+    coros = []
+    with flyte.group("fanout-group"):
+        for id in ids:
+            coros.append(hello_worker(id))
+
+        vals = await tqdm.gather(*coros,desc="Running Tasks")
+
+    return vals
+
+
+if __name__ == "__main__":
+    flyte.init_from_config()
+
+    run = flyte.run(hello_driver)
+    print(run.name)
+    print(run.url)
+    run.wait()

maint_tools/build_default_image.py

@@ -1,5 +1,6 @@
 import argparse
 import asyncio
+from tqdm.asyncio import tqdm
 from pathlib import Path
 
 import flyte
@@ -72,9 +73,10 @@ async def build_flyte_connector_image(
 
 
 async def build_all(registry: str | None = None, name: str | None = None, builder: str | None = "local"):
-    await asyncio.gather(
+    await tqdm.gather(
         build_flyte_image(registry=registry, name=name, builder=builder),
         build_flyte_connector_image(registry=registry, name=name, builder=builder),
+        desc="Building Images"
     )
 
 

maint_tools/build_plugin_image.py

@@ -1,6 +1,6 @@
 import argparse
 import asyncio
-
+from tqdm.asyncio import tqdm
 import flyte
 from flyte.extend import ImageBuildEngine
 
@@ -33,9 +33,10 @@ async def build_flyte_sglang_image(registry: str | None = None, name: str | None
 
 
 async def build_all(registry: str | None = None, name: str | None = None, builder: str | None = "local"):
-    await asyncio.gather(
+    await tqdm.gather(
         build_flyte_vllm_image(registry=registry, name=name, builder=builder),
         build_flyte_sglang_image(registry=registry, name=name, builder=builder),
+        desc="Building plugin images"
     )
 
 

pyproject.toml

@@ -26,6 +26,7 @@ dependencies = [
     "aiolimiter>=1.2.1",
     "flyteidl2==2.0.2",
     "packaging",
+    "tqdm>=4.67.3",
 ]
 
 [project.scripts]

src/flyte/_deploy.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import asyncio
+from tqdm.asyncio import tqdm
 import hashlib
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Dict, List, Optional, Set, Tuple
@@ -396,7 +397,7 @@ async def _build_images(deployment: DeploymentPlan, image_refs: Dict[str, str] |
                 continue
             auto_image = Image.from_debian_base()
             images.append(_build_image_bg(env_name, auto_image))
-    final_images = await asyncio.gather(*images)
+    final_images = await tqdm.gather(*images,desc="Building Images")
 
     for env_name, image_uri in final_images:
         logger.warning(f"Built Image for environment {env_name}, image: {image_uri}")
@@ -417,7 +418,7 @@ async def _deploy_task_env(context: DeploymentContext) -> DeployedTaskEnvironmen
     task_coros = []
     for task in env.tasks.values():
         task_coros.append(_deploy_task(task, context.serialization_context, dryrun=context.dryrun))
-    deployed_task_vals = await asyncio.gather(*task_coros)
+    deployed_task_vals = await tqdm.gather(*task_coros,desc="Deploying Task Environment")
     deployed_tasks = []
     for t in deployed_task_vals:
         deployed_tasks.append(t)
@@ -466,7 +467,7 @@ async def apply(deployment_plan: DeploymentPlan, copy_style: CopyFiles, dryrun:
         deployer = get_deployer(type(env))
         context = DeploymentContext(environment=env, serialization_context=sc, dryrun=dryrun)
         deployment_coros.append(deployer(context))
-    deployed_envs = await asyncio.gather(*deployment_coros)
+    deployed_envs = await tqdm.gather(*deployment_coros,desc="Deploying Environment")
     envs = {}
     for d in deployed_envs:
         envs[d.get_name()] = d
@@ -533,7 +534,7 @@ async def deploy(
     deployments = []
     for deployment_plan in deployment_plans:
         deployments.append(apply(deployment_plan, copy_style=copy_style, dryrun=dryrun))
-    return await asyncio.gather(*deployments)
+    return await tqdm.gather(*deployments,desc="Deployment in progress")
 
 
 @syncify

src/flyte/_serve.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import asyncio
+from tqdm.asyncio import tqdm
 import atexit
 import hashlib
 import os
@@ -663,7 +664,7 @@ async def _serve_remote(self, app_env: "AppEnvironment") -> "App":
                 app_envs_to_deploy.append(dep_env)
 
         # Deploy all apps concurrently
-        deployed_apps = await asyncio.gather(*deployment_coros)
+        deployed_apps = await tqdm.gather(*deployment_coros,desc=f"Deploying Apps")
 
         # Find the deployed app corresponding to the requested app_env
         deployed_app = None

src/flyte/remote/_data.py

@@ -1,4 +1,5 @@
 import asyncio
+from tqdm.asyncio import tqdm
 import hashlib
 import os
 import typing
@@ -234,7 +235,7 @@ async def upload_dir(dir_path: Path, verify: bool = True, prefix: str | None = N
         if file.is_file():
             uploaded_files.append(_upload_single_file(cfg, file, verify=verify, basedir=prefix))
 
-    urls = await asyncio.gather(*uploaded_files)
+    urls = await tqdm.gather(*uploaded_files,desc=f"Uploading files to {dir_path}")
     native_url = urls[0][1]  # Assuming all files are uploaded to the same prefix
     # native_url is of the form s3://my-s3-bucket/flytesnacks/development/{prefix}/source/empty.md
     uri = native_url.split(prefix)[0]