feat(tidy3d): FXC-3886-cache-logging-verbosity-control

Author: marcorudolphflex

tests/test_web/test_local_cache.py

@@ -1,20 +1,26 @@
 from __future__ import annotations
 
+import importlib
+import io
 import os
+import re
 from pathlib import Path
+from types import SimpleNamespace
 
 import pytest
+from rich.console import Console
 
 import tidy3d as td
 from tests.test_components.autograd.test_autograd import ALL_KEY, get_functions, params0
 from tests.test_web.test_webapi_mode import make_mode_sim
 from tidy3d import config
 from tidy3d.config import get_manager
-from tidy3d.web import Job, common, run_async
+from tidy3d.web import Job, common, run, run_async
 from tidy3d.web.api import webapi as web
 from tidy3d.web.api.container import Batch, WebContainer
 from tidy3d.web.api.webapi import load_simulation_if_cached
 from tidy3d.web.cache import CACHE_ARTIFACT_NAME, clear, get_cache_entry_dir, resolve_local_cache
+from tidy3d.web.core.task_core import BatchTask
 
 common.CONNECTION_RETRY_TIME = 0.1
 
@@ -139,6 +145,8 @@ def _fake_status(self):
     monkeypatch.setattr(
         web, "load_simulation", lambda task_id, *args, **kwargs: TASK_TO_SIM[task_id]
     )
+    monkeypatch.setattr(BatchTask, "is_batch", lambda *args, **kwargs: "success")
+    monkeypatch.setattr(BatchTask, "detail", lambda *args: SimpleNamespace(status="success"))
     return counters
 
 
@@ -267,6 +275,79 @@ def _test_run_cache_hit_async(monkeypatch, basic_simulation, tmp_path):
     assert len(cache) == 3
 
 
+def _test_verbosity(monkeypatch, basic_simulation):
+    _CSI_RE = re.compile(r"\x1b\[[0-?]*[ -/]*[@-~]")  # ANSI CSI
+    _OSC8_RE = re.compile(r"\x1b\]8;.*?(?:\x1b\\|\x07)", re.DOTALL)  # OSC-8 hyperlinks
+
+    def _normalize_console_text(s: str) -> str:
+        s = _OSC8_RE.sub("", s)  # drop hyperlinks
+        s = _CSI_RE.sub("", s)  # drop colors, etc.
+        s = s.replace("\r", "")
+        s = re.sub(r"[ \t]+", " ", s)  # collapse runs of spaces/tabs
+        return s.strip()
+
+    buf = io.StringIO()
+    test_console = Console(file=buf, force_terminal=True)
+
+    counters = _patch_run_pipeline(monkeypatch)
+    cache = resolve_local_cache(use_cache=True)
+    cache.clear()
+    _reset_fake_maps()
+    _reset_counters(counters)
+    sim2 = basic_simulation.updated_copy(shutoff=1e-4)
+    sim3 = basic_simulation.updated_copy(shutoff=1e-3)
+
+    run(basic_simulation, verbose=True)  # seed cache
+
+    log_mod = importlib.import_module("tidy3d.log")
+
+    # --- swap handler console (and restore later) ---
+    if "console" not in log_mod.log.handlers:
+        log_mod.set_logging_console()
+    orig_console = log_mod.log.handlers["console"].console
+    log_mod.log.handlers["console"].console = test_console
+    try:
+        buf.truncate(0)
+        buf.seek(0)
+
+        log_mod.get_logging_console().log("test")
+        assert "test" in buf.getvalue()
+
+        buf.truncate(0)
+        buf.seek(0)
+
+        # test for load_simulation_if_cached
+        sim_data = load_simulation_if_cached(basic_simulation, verbose=True)
+        assert sim_data is not None
+        assert "Loading simulation from" in buf.getvalue(), (
+            f"Expected 'Loading simulation from' in log, got '{buf.getvalue()}'"
+        )
+
+        buf.truncate(0)
+        buf.seek(0)
+        load_simulation_if_cached(basic_simulation, verbose=False)
+        assert sim_data is not None
+        assert buf.getvalue().strip() == "", f"Expected empty log, got '{buf.getvalue()}'"
+
+        # test for batched runs
+        buf.truncate(0)
+        buf.seek(0)
+        run([basic_simulation, sim3], verbose=True)
+        txt = _normalize_console_text(buf.getvalue())
+        assert "Got 1 simulation from cache" in txt, (
+            f"Expected 'Got 1 simulation from cache' in log, got '{buf.getvalue()}'"
+        )
+
+        buf.truncate(0)
+        buf.seek(0)
+        run([basic_simulation, sim2], verbose=False)
+        assert buf.getvalue().strip() == "", f"Expected empty log, got '{buf.getvalue()}'"
+
+    finally:
+        # explicit restore so nothing leaks
+        log_mod.log.handlers["console"].console = orig_console
+
+
 def _test_job_run_cache(monkeypatch, basic_simulation, tmp_path):
     counters = _patch_run_pipeline(monkeypatch)
     cache = resolve_local_cache(use_cache=True)
@@ -436,3 +517,4 @@ def test_cache_sequential(monkeypatch, tmp_path, tmp_path_factory, basic_simulat
     _test_autograd_cache(monkeypatch)
     _test_configure_cache_roundtrip(monkeypatch, tmp_path)
     _test_mode_solver_caching(monkeypatch, tmp_path)
+    _test_verbosity(monkeypatch, basic_simulation)

tidy3d/web/api/container.py

@@ -18,7 +18,14 @@
 
 import pydantic.v1 as pd
 from pydantic.v1 import PrivateAttr
-from rich.progress import BarColumn, Progress, TaskProgressColumn, TextColumn, TimeElapsedColumn
+from rich.progress import (
+    BarColumn,
+    Progress,
+    TaskID,
+    TaskProgressColumn,
+    TextColumn,
+    TimeElapsedColumn,
+)
 
 from tidy3d.components.base import Tidy3dBaseModel, cached_property
 from tidy3d.components.mode.mode_solver import ModeSolver
@@ -331,7 +338,7 @@ def load_if_cached(self) -> bool:
             simulation=self.simulation,
             path=stash_path,
             reduce_simulation=self.reduce_simulation,
-            verbose=getattr(self, "verbose", True),
+            verbose=self.verbose,
         )
 
         if restored is None:
@@ -945,27 +952,42 @@ def upload(self) -> None:
         """Upload a series of tasks associated with this ``Batch`` using multi-threading."""
         self._check_folder(self.folder_name)
         with ThreadPoolExecutor(max_workers=self.num_workers) as executor:
-            futures = [executor.submit(job.upload) for _, job in self.jobs.items()]
+            jobs_from_cache = [job for job in self.jobs.values() if job.load_if_cached]
+            jobs_to_upload = [job for job in self.jobs.values() if not job.load_if_cached]
+            futures = [executor.submit(job.upload) for job in jobs_to_upload]
 
             # progressbar (number of tasks uploaded)
             if self.verbose:
                 console = get_logging_console()
+                n_cached = len(jobs_from_cache)
+                if n_cached > 0:
+                    console.log(
+                        f"Got {n_cached} simulation{'s' if n_cached > 1 else ''} from cache."
+                    )
+
+            if len(futures) == 0:  # got all jobs from cache
+                return
+
+            if self.verbose:
                 progress_columns = (
                     TextColumn("[progress.description]{task.description}"),
                     BarColumn(),
                     TaskProgressColumn(),
                     TimeElapsedColumn(),
                 )
                 with Progress(*progress_columns, console=console, transient=False) as progress:
-                    pbar_message = f"Uploading data for {self.num_jobs} tasks"
-                    pbar = progress.add_task(pbar_message, total=self.num_jobs)
+                    pbar_message = f"Uploading data for {len(jobs_to_upload)} task{'s' if len(jobs_to_upload) > 1 else ''}"
+                    pbar = progress.add_task(pbar_message, total=len(jobs_to_upload))
                     completed = 0
                     for _ in concurrent.futures.as_completed(futures):
                         completed += 1
                         progress.update(pbar, completed=completed)
 
                     progress.refresh()
                     time.sleep(BATCH_PROGRESS_REFRESH_TIME)
+            else:
+                for _ in concurrent.futures.as_completed(futures):
+                    pass
 
     def get_info(self) -> dict[TaskName, TaskInfo]:
         """Get information about each task in the :class:`Batch`.
@@ -1138,6 +1160,8 @@ def pbar_description(
         postprocess_started_tasks: set[str] = set()
 
         try:
+            console = None
+            progress_columns = []
             if self.verbose:
                 console = get_logging_console()
                 self.estimate_cost()
@@ -1152,34 +1176,38 @@ def pbar_description(
                     TimeElapsedColumn(),
                 )
 
-                with Progress(*progress_columns, console=console, transient=False) as progress:
-                    pbar_tasks: dict[str, int] = {}
-                    for task_name, job in self.jobs.items():
-                        schedule_download(job)
+            with Progress(
+                *progress_columns, console=console, transient=False, disable=not self.verbose
+            ) as progress:
+                pbar_tasks: dict[str, TaskID] = {}
+                for task_name, job in self.jobs.items():
+                    schedule_download(job)
+                    if self.verbose:
                         status = job.status
                         completed = STATE_PROGRESS_PERCENTAGE.get(status, 0)
                         desc = pbar_description(task_name, status, max_name_length, 0)
                         pbar_tasks[task_name] = progress.add_task(
                             desc, total=COMPLETED_PERCENT, completed=completed
                         )
 
-                    while any(check_continue_condition(job) for job in self.jobs.values()):
-                        for task_name, job in self.jobs.items():
-                            status = job.status
-
-                            # auto-start postprocess for modeler jobs when run finishes
-                            if (
-                                web._is_modeler_batch(job.task_id)
-                                and status == "run_success"
-                                and job.task_id not in postprocess_started_tasks
-                            ):
-                                job.postprocess_start(
-                                    worker_group=postprocess_worker_group, verbose=True
-                                )
-                                postprocess_started_tasks.add(job.task_id)
+                while any(check_continue_condition(job) for job in self.jobs.values()):
+                    for task_name, job in self.jobs.items():
+                        status = job.status
+
+                        # auto-start postprocess for modeler jobs when run finishes
+                        if (
+                            web._is_modeler_batch(job.task_id)
+                            and status == "run_success"
+                            and job.task_id not in postprocess_started_tasks
+                        ):
+                            job.postprocess_start(
+                                worker_group=postprocess_worker_group, verbose=True
+                            )
+                            postprocess_started_tasks.add(job.task_id)
 
-                            schedule_download(job)
+                        schedule_download(job)
 
+                        if self.verbose:
                             # choose display status & percent
                             if status != "run_success":
                                 display_status = status
@@ -1196,13 +1224,17 @@ def pbar_description(
                             pbar = pbar_tasks[task_name]
                             desc = pbar_description(task_name, display_status, max_name_length, 0)
                             progress.update(pbar, description=desc, completed=pct)
-
+                    if self.verbose:
                         progress.refresh()
                         time.sleep(BATCH_PROGRESS_REFRESH_TIME)
+                    else:
+                        time.sleep(web.REFRESH_TIME)
 
-                    # final render to terminal state for all bars
-                    for task_name, job in self.jobs.items():
-                        schedule_download(job)
+                # final render to terminal state for all bars
+                for task_name, job in self.jobs.items():
+                    schedule_download(job)
+
+                    if self.verbose:
                         status = job.status
                         if status != "run_success":
                             display_status = status
@@ -1222,25 +1254,9 @@ def pbar_description(
                         desc = pbar_description(task_name, display_status, max_name_length, 0)
                         progress.update(pbar, description=desc, completed=pct)
 
+                if self.verbose:
                     progress.refresh()
                     console.log("Batch complete.")
-            else:
-                # quiet mode
-                while any(check_continue_condition(job) for job in self.jobs.values()):
-                    for job in self.jobs.values():
-                        if (
-                            web._is_modeler_batch(job.task_id)
-                            and job.status == "run_success"
-                            and job.task_id not in postprocess_started_tasks
-                        ):
-                            job.postprocess_start(
-                                worker_group=postprocess_worker_group, verbose=False
-                            )
-                            postprocess_started_tasks.add(job.task_id)
-
-                        schedule_download(job)
-
-                    time.sleep(web.REFRESH_TIME)
         finally:
             if download_executor is not None:
                 try:

tidy3d/web/api/webapi.py

@@ -387,10 +387,10 @@ def restore_simulation_if_cached(
             cached_task_id = entry.metadata.get("task_id")
             cached_workflow_type = entry.metadata.get("workflow_type")
             if cached_task_id is not None and cached_workflow_type is not None and verbose:
-                console = get_logging_console() if verbose else None
+                console = get_logging_console()
                 url, _ = _get_task_urls(cached_workflow_type, simulation, cached_task_id)
                 console.log(
-                    f"Loaded simulation from local cache.\nView cached task using web UI at [link={url}]'{url}'[/link]."
+                    f"Loading simulation from local cache. View cached task using web UI at [link={url}]'{url}'[/link]."
                 )
     return retrieved_simulation_path
 
@@ -399,6 +399,7 @@ def load_simulation_if_cached(
     simulation: WorkflowType,
     path: Optional[PathLike] = None,
     reduce_simulation: Literal["auto", True, False] = "auto",
+    verbose: bool = True,
 ) -> Optional[WorkflowDataType]:
     """
     Load simulation results directly from the local cache, if available.
@@ -412,17 +413,22 @@ def load_simulation_if_cached(
     reduce_simulation : Literal["auto", True, False] = "auto"
         Whether to use a reduced simulation when checking the cache. If "auto",
         reduction is applied automatically for mode solvers.
+    verbose : bool = True
+        If True, logs a message including a link to the cached task in the web UI on loading.
 
     Returns
     -------
     Optional[WorkflowDataType]
         The loaded simulation data if found in cache, otherwise None.
     """
-    restored_path = restore_simulation_if_cached(simulation, path, reduce_simulation)
+    restored_path = restore_simulation_if_cached(
+        simulation, path, reduce_simulation, verbose=verbose
+    )
     if restored_path is not None:
         data = load(
             task_id=None,
             path=str(restored_path),
+            verbose=verbose,
         )
         if isinstance(simulation, ModeSolver):
             simulation._patch_data(data=data)