From 4e6852761081388a316e0af58b1cede0539d0365 Mon Sep 17 00:00:00 2001 From: Wu-Jung Lee Date: Fri, 14 Aug 2026 10:47:27 -0700 Subject: [PATCH 1/8] change flow_raw2Sv_postprocessing to be based off local raw files --- src/echodataflow/flows/flows_acoustics.py | 24 ++++++++++------------- 1 file changed, 10 insertions(+), 14 deletions(-) diff --git a/src/echodataflow/flows/flows_acoustics.py b/src/echodataflow/flows/flows_acoustics.py index 6068c4f7..3f11c1eb 100644 --- a/src/echodataflow/flows/flows_acoustics.py +++ b/src/echodataflow/flows/flows_acoustics.py @@ -38,12 +38,10 @@ plan_mvbs_slices, read_or_create_ledger, ) -from echodataflow.operations.operations_storage import S3CopySettings, S3CopyWorkItem from echodataflow.tasks.tasks_acoustics import ( task_create_MVBS, task_raw2Sv, ) -from echodataflow.tasks.tasks_postprocessing import task_s3_raw2Sv from echodataflow.utils.utils import ( round_up_mins, get_slice_start_end_times, @@ -408,9 +406,8 @@ async def flow_create_MVBS( @flow(log_prints=True, task_runner=dask_task_runner_from_environment()) def flow_raw2Sv_postprocessing( path_raw_list: str, + path_raw: str, path_main: str, - s3_bucket: str = "noaa-wcsd-pds", - endpoint_url: str | None = "https://sdsc.osn.xsede.org", start_time: str | None = None, end_time: str | None = None, new_file_num_limit: int = -1, @@ -424,12 +421,13 @@ def flow_raw2Sv_postprocessing( nmea_sentence: str | None = None, file_Sv_csv: str = "Sv_files.csv", ) -> None: - """Convert raw files and update corresponding rows in the Sv ledger.""" + """Convert local raw files and update corresponding rows in the Sv ledger.""" logger = get_run_logger() - # Keep temporary raw files separate from persistent Sv outputs - path_raw_staging = Path(path_main) / "raw_staging" + path_raw = Path(path_raw) + if not path_raw.is_dir(): + raise ValueError(f"Local raw directory does not exist: {path_raw}") + path_Sv = Path(path_main) / "Sv" - path_raw_staging.mkdir(parents=True, exist_ok=True) path_Sv.mkdir(parents=True, exist_ok=True) file_Sv_csv = Path(path_main) / file_Sv_csv @@ -458,7 +456,6 @@ def flow_raw2Sv_postprocessing( df_Sv.loc[selected.index, ["raw2Sv_status", "error"]] = ["pending", ""] write_manifest(df_Sv, file_Sv_csv) - copy_settings = S3CopySettings(s3_bucket=s3_bucket, endpoint_url=endpoint_url) sv_settings = RawToSvSettings( output_directory=str(path_Sv), encode_mode=encode_mode, @@ -469,19 +466,18 @@ def flow_raw2Sv_postprocessing( nmea_sentence=nmea_sentence, ) - # Submit one download-plus-conversion task per raw object + # Submit one conversion task per local raw file errors = [] conversion_futures = {} for row in selected.itertuples(index=False): key = str(row.s3_path) filename = Path(key).name - future = task_s3_raw2Sv.with_options( + future = task_raw2Sv.with_options( task_run_name=f"raw2Sv_{filename}", retries=task_retries, retry_delay_seconds=task_retry_delay_seconds, ).submit( - S3CopyWorkItem(s3_path=key, local_path=str(path_raw_staging / filename)), - copy_settings, + RawToSvWorkItem(raw_path=str(path_raw / filename)), sv_settings, ) conversion_futures[future] = key @@ -516,7 +512,7 @@ def flow_raw2Sv_postprocessing( errors.append(exc) df_Sv.loc[idx, ["raw2Sv_status", "error"]] = ["failed", str(exc)] write_manifest(df_Sv, file_Sv_csv) - logger.error("Failed to download or convert %s: %s", key, exc) + logger.error("Failed to convert %s: %s", key, exc) if errors: raise RuntimeError(f"{len(errors)} raw-to-Sv conversions failed") From 6979dafe8c9a643c87b2de70d7f53a83b9b1747d Mon Sep 17 00:00:00 2001 From: Wu-Jung Lee Date: Fri, 14 Aug 2026 11:01:47 -0700 Subject: [PATCH 2/8] consolidate cancel is deployment already running and add to all postproc flows --- src/echodataflow/flows/flows_acoustics.py | 23 ++++------ src/echodataflow/flows/flows_helper.py | 18 ++++++++ src/echodataflow/flows/flows_predict_hake.py | 5 ++ tests/test_postprocessing_flow_validation.py | 48 ++++++++++++++++++++ 4 files changed, 79 insertions(+), 15 deletions(-) diff --git a/src/echodataflow/flows/flows_acoustics.py b/src/echodataflow/flows/flows_acoustics.py index 3f11c1eb..f283692f 100644 --- a/src/echodataflow/flows/flows_acoustics.py +++ b/src/echodataflow/flows/flows_acoustics.py @@ -10,10 +10,10 @@ from prefect import flow, get_run_logger, get_client from prefect.futures import as_completed -from prefect.states import Cancelled, Failed +from prefect.states import Failed from prefect import runtime -from echodataflow.flows.flows_helper import deployment_already_running +from echodataflow.flows.flows_helper import cancel_if_deployment_already_running from echodataflow.deployment.task_runners import dask_task_runner_from_environment from echodataflow.utils.manifests import ( MVBS_COLUMNS_POSTPROCESSING, @@ -71,19 +71,8 @@ def flow_raw2Sv( new_file_num_limit: int = 50, ): - # Check if the deployment is already running - already_running = asyncio.run(deployment_already_running()) - if already_running: - - async def cancel_run(): - async with get_client() as client: - await client.set_flow_run_state( - flow_run_id=runtime.flow_run.id, - state=Cancelled(message="Another instance of this flow is already running"), - ) - - asyncio.run(cancel_run()) - return # exit the flow early + if cancel_if_deployment_already_running(): + return # Assemble paths path_Sv_zarr = Path(path_main) / "Sv" @@ -527,6 +516,10 @@ def flow_create_MVBS_postprocessing( file_MVBS_csv: str = "MVBS_files.csv", ) -> None: """Create preplanned MVBS slices after all required raw conversions finish.""" + + if cancel_if_deployment_already_running(): + return + logger = get_run_logger() file_Sv_csv = Path(path_main) / file_Sv_csv if not file_Sv_csv.exists(): diff --git a/src/echodataflow/flows/flows_helper.py b/src/echodataflow/flows/flows_helper.py index b8a3064c..a6a2670a 100644 --- a/src/echodataflow/flows/flows_helper.py +++ b/src/echodataflow/flows/flows_helper.py @@ -1,8 +1,10 @@ from pathlib import Path +import asyncio import datetime from prefect import flow, get_client, runtime, task from prefect.client.schemas.filters import FlowRunFilter +from prefect.states import Cancelled @flow(timeout_seconds=600, log_prints=True) @@ -92,3 +94,19 @@ async def deployment_already_running() -> bool: ) return len(running_flows) > 1 + + +def cancel_if_deployment_already_running() -> bool: + """Cancel this flow run and return True when its deployment already has a run.""" + if not asyncio.run(deployment_already_running()): + return False + + async def cancel_run(): + async with get_client() as client: + await client.set_flow_run_state( + flow_run_id=runtime.flow_run.id, + state=Cancelled(message="Another instance of this flow is already running"), + ) + + asyncio.run(cancel_run()) + return True diff --git a/src/echodataflow/flows/flows_predict_hake.py b/src/echodataflow/flows/flows_predict_hake.py index 5ac39548..db750171 100644 --- a/src/echodataflow/flows/flows_predict_hake.py +++ b/src/echodataflow/flows/flows_predict_hake.py @@ -10,6 +10,7 @@ from prefect import flow, get_client, get_run_logger, runtime from prefect.states import Failed +from echodataflow.flows.flows_helper import cancel_if_deployment_already_running from echodataflow.utils.manifests import ( MVBS_COLUMNS_POSTPROCESSING, PREDICTION_COLUMNS_POSTPROCESSING, @@ -260,6 +261,10 @@ def flow_predict_hake_postprocessing( file_prediction_csv: str = "prediction_files.csv", ) -> None: """Predict all newly ready windows, combining aligned MVBS slices.""" + + if cancel_if_deployment_already_running(): + return + logger = get_run_logger() file_MVBS_csv = Path(path_main) / file_MVBS_csv if not file_MVBS_csv.exists(): diff --git a/tests/test_postprocessing_flow_validation.py b/tests/test_postprocessing_flow_validation.py index bac5c645..a480df0d 100644 --- a/tests/test_postprocessing_flow_validation.py +++ b/tests/test_postprocessing_flow_validation.py @@ -3,6 +3,7 @@ from echodataflow.flows.flows_acoustics import ( flow_create_MVBS_postprocessing, + flow_raw2Sv, ) from echodataflow.flows.flows_predict_hake import flow_predict_hake_postprocessing @@ -15,6 +16,11 @@ def info(self, message): messages.append(message) monkeypatch.setattr(flows_acoustics, "get_run_logger", lambda: Logger()) + monkeypatch.setattr( + flows_acoustics, + "cancel_if_deployment_already_running", + lambda: False, + ) flow_create_MVBS_postprocessing.fn(path_main=str(tmp_path)) @@ -30,6 +36,11 @@ def info(self, message): messages.append(message) monkeypatch.setattr(flows_predict_hake, "get_run_logger", lambda: Logger()) + monkeypatch.setattr( + flows_predict_hake, + "cancel_if_deployment_already_running", + lambda: False, + ) flow_predict_hake_postprocessing.fn( path_main=str(tmp_path), @@ -38,3 +49,40 @@ def info(self, message): assert messages == ["MVBS ledger does not yet exist"] assert not (tmp_path / "prediction_files.csv").exists() + + +def test_mvbs_flow_exits_when_deployment_is_already_running(tmp_path, monkeypatch): + monkeypatch.setattr( + flows_acoustics, + "cancel_if_deployment_already_running", + lambda: True, + ) + + flow_create_MVBS_postprocessing.fn(path_main=str(tmp_path)) + + assert not (tmp_path / "MVBS_files.csv").exists() + + +def test_realtime_raw2sv_flow_exits_when_deployment_is_already_running(monkeypatch): + monkeypatch.setattr( + flows_acoustics, + "cancel_if_deployment_already_running", + lambda: True, + ) + + assert flow_raw2Sv.fn() is None + + +def test_prediction_flow_exits_when_deployment_is_already_running(tmp_path, monkeypatch): + monkeypatch.setattr( + flows_predict_hake, + "cancel_if_deployment_already_running", + lambda: True, + ) + + flow_predict_hake_postprocessing.fn( + path_main=str(tmp_path), + path_weight="unused.ckpt", + ) + + assert not (tmp_path / "prediction_files.csv").exists() From 8542164a1e4ef0c3603ec12bc4999906c6812dbb Mon Sep 17 00:00:00 2001 From: Wu-Jung Lee Date: Fri, 14 Aug 2026 11:35:28 -0700 Subject: [PATCH 3/8] use prefect native ConcurrencyLimitStrategy to cancel new flow when one prev flow of the same deployment is already running --- src/echodataflow/deployment/core.py | 2 + .../deployment/deployment_engine.py | 45 +++++++++++++--- src/echodataflow/flows/flows_acoustics.py | 7 --- src/echodataflow/flows/flows_helper.py | 41 +-------------- src/echodataflow/flows/flows_predict_hake.py | 4 -- tests/deployment/conftest.py | 23 ++++++++- tests/deployment/test_deploy_engine.py | 51 +++++++++++++++++++ tests/test_postprocessing_flow_validation.py | 50 ------------------ 8 files changed, 115 insertions(+), 108 deletions(-) diff --git a/src/echodataflow/deployment/core.py b/src/echodataflow/deployment/core.py index 8fe9301f..b4509565 100644 --- a/src/echodataflow/deployment/core.py +++ b/src/echodataflow/deployment/core.py @@ -14,12 +14,14 @@ "flow", "interval", "cron_offset", + "concurrency_limit", "triggers", "inject_time_offset", "task_runner", "work_pool_name", } ALLOWED_CONCURRENCY_GROUP_KEYS = {"limit"} +ALLOWED_CONCURRENCY_LIMIT_KEYS = {"limit", "collision_strategy", "grace_period_seconds"} ALLOWED_TASK_RUNNER_KEYS = {"type", "cluster_kwargs"} ALLOWED_DASK_CLUSTER_KEYS = { "memory_limit", diff --git a/src/echodataflow/deployment/deployment_engine.py b/src/echodataflow/deployment/deployment_engine.py index 40b3b662..2ee1026d 100644 --- a/src/echodataflow/deployment/deployment_engine.py +++ b/src/echodataflow/deployment/deployment_engine.py @@ -18,6 +18,7 @@ from echodataflow.deployment.core import ( ALLOWED_CONCURRENCY_GROUP_KEYS, + ALLOWED_CONCURRENCY_LIMIT_KEYS, ALLOWED_DASK_CLUSTER_KEYS, ALLOWED_DEPLOY_KEYS, ALLOWED_FLOW_DEPLOY_KEYS, @@ -40,6 +41,7 @@ class DeploymentSpec: entrypoint: str # source-relative entrypoint for the actual deployed flow parameters: dict[str, Any] # parameters passed directly to the deployed flow concurrency_group: str | None = None + concurrency_limit: dict[str, Any] | None = None task_runner: dict[str, Any] | None = None cron: str | None = None # precomputed cron schedule, when interval mode is used work_pool_name: str | None = ( @@ -339,6 +341,25 @@ def validate_deploy_config(deploy_cfg: Any) -> None: if task_runner is not None: validate_task_runner_config(task_runner, path=f"{flow_path}.task_runner") + concurrency_limit = deploy_meta.get("concurrency_limit") + if concurrency_limit is not None: + limit_path = f"{flow_path}.concurrency_limit" + if not isinstance(concurrency_limit, dict): + raise ValueError(f"{limit_path} must be a mapping") + _reject_unknown_keys( + concurrency_limit, + allowed=ALLOWED_CONCURRENCY_LIMIT_KEYS, + path=limit_path, + ) + if "limit" not in concurrency_limit: + raise ValueError(f"{limit_path}.limit is required") + _validate_positive_integer(concurrency_limit["limit"], path=f"{limit_path}.limit") + strategy = concurrency_limit.get("collision_strategy", "ENQUEUE") + if strategy not in {"ENQUEUE", "CANCEL_NEW"}: + raise ValueError( + f"{limit_path}.collision_strategy must be 'ENQUEUE' or 'CANCEL_NEW'" + ) + triggers = deploy_meta.get("triggers") if isinstance(triggers, list): for index, trigger in enumerate(triggers): @@ -457,8 +478,8 @@ def validate_flow_coverage( raise ValueError("Flow coverage mismatch. " + " | ".join(errors)) -def _flow_accepts_time_offset_seconds(flow_obj: Any) -> bool: - """Return True when the flow function can accept `time_offset_seconds`. +def _flow_accepts_parameter(flow_obj: Any, parameter: str) -> bool: + """Return True when the flow function accepts the named parameter. Prefect Flow objects expose the wrapped function via `.fn`. If a flow object does not expose an inspectable function (e.g. certain test doubles), we skip @@ -469,7 +490,7 @@ def _flow_accepts_time_offset_seconds(flow_obj: Any) -> bool: return True signature = inspect.signature(flow_fn) - return "time_offset_seconds" in signature.parameters + return parameter in signature.parameters def build_deploy_specs( @@ -496,8 +517,8 @@ def build_deploy_specs( flow_info = resolved_flows[key] # Check if time_offset_seconds is indeed accepted by the flows specified in deploy config - if key in time_offset_targets and not _flow_accepts_time_offset_seconds( - flow_info["flow_obj"] + if key in time_offset_targets and not _flow_accepts_parameter( + flow_info["flow_obj"], "time_offset_seconds" ): raise ValueError( f"deploy_cfg.flows.{key}.inject_time_offset is enabled, " @@ -531,7 +552,6 @@ def build_deploy_specs( deployment_parameters = dict(flow_params) if key in time_offset_targets: deployment_parameters["time_offset_seconds"] = time_offset_seconds - specs.append( DeploymentSpec( flow_key=key, @@ -540,6 +560,7 @@ def build_deploy_specs( entrypoint=flow_info["entrypoint"], parameters=deployment_parameters, concurrency_group=deploy_meta.get("concurrency_group"), + concurrency_limit=deploy_meta.get("concurrency_limit"), task_runner=deploy_meta.get("task_runner"), cron=cron, work_pool_name=deploy_meta.get("work_pool_name"), @@ -582,6 +603,18 @@ def create_deployments( if spec.concurrency_group is not None: deployment_kwargs["work_queue_name"] = spec.concurrency_group + if spec.concurrency_limit is not None: + from prefect.client.schemas.objects import ( + ConcurrencyLimitConfig, + ConcurrencyLimitStrategy, + ) + + concurrency_limit = dict(spec.concurrency_limit) + concurrency_limit["collision_strategy"] = ConcurrencyLimitStrategy( + concurrency_limit.get("collision_strategy", "ENQUEUE") + ) + deployment_kwargs["concurrency_limit"] = ConcurrencyLimitConfig(**concurrency_limit) + # The worker reloads the flow entrypoint, so runner settings must be # present in its runtime environment instead of only on this Flow object if spec.task_runner is not None: diff --git a/src/echodataflow/flows/flows_acoustics.py b/src/echodataflow/flows/flows_acoustics.py index f283692f..3a7e7bd8 100644 --- a/src/echodataflow/flows/flows_acoustics.py +++ b/src/echodataflow/flows/flows_acoustics.py @@ -13,7 +13,6 @@ from prefect.states import Failed from prefect import runtime -from echodataflow.flows.flows_helper import cancel_if_deployment_already_running from echodataflow.deployment.task_runners import dask_task_runner_from_environment from echodataflow.utils.manifests import ( MVBS_COLUMNS_POSTPROCESSING, @@ -71,9 +70,6 @@ def flow_raw2Sv( new_file_num_limit: int = 50, ): - if cancel_if_deployment_already_running(): - return - # Assemble paths path_Sv_zarr = Path(path_main) / "Sv" file_Sv_csv = Path(path_main) / file_Sv_csv @@ -516,9 +512,6 @@ def flow_create_MVBS_postprocessing( file_MVBS_csv: str = "MVBS_files.csv", ) -> None: """Create preplanned MVBS slices after all required raw conversions finish.""" - - if cancel_if_deployment_already_running(): - return logger = get_run_logger() file_Sv_csv = Path(path_main) / file_Sv_csv diff --git a/src/echodataflow/flows/flows_helper.py b/src/echodataflow/flows/flows_helper.py index a6a2670a..b4c57ab4 100644 --- a/src/echodataflow/flows/flows_helper.py +++ b/src/echodataflow/flows/flows_helper.py @@ -1,10 +1,7 @@ from pathlib import Path -import asyncio import datetime -from prefect import flow, get_client, runtime, task -from prefect.client.schemas.filters import FlowRunFilter -from prefect.states import Cancelled +from prefect import flow @flow(timeout_seconds=600, log_prints=True) @@ -74,39 +71,3 @@ def flow_file_upload( # Remove the exclude list file after upload exclude_path.unlink(missing_ok=True) - - -@task(log_prints=True) -async def deployment_already_running() -> bool: - """Return whether another run of the current deployment is running.""" - # Not running as a deployment, so skip the check - if runtime.deployment.id is None: - return False - - # Check if the deployment is already running - async with get_client() as client: - # Get all running flows for this deployment using simpler filters - running_flows = await client.read_flow_runs( - flow_run_filter=FlowRunFilter( - deployment_id={"any_": [runtime.deployment.id]}, - state={"type": {"any_": ["RUNNING"]}}, - ) - ) - - return len(running_flows) > 1 - - -def cancel_if_deployment_already_running() -> bool: - """Cancel this flow run and return True when its deployment already has a run.""" - if not asyncio.run(deployment_already_running()): - return False - - async def cancel_run(): - async with get_client() as client: - await client.set_flow_run_state( - flow_run_id=runtime.flow_run.id, - state=Cancelled(message="Another instance of this flow is already running"), - ) - - asyncio.run(cancel_run()) - return True diff --git a/src/echodataflow/flows/flows_predict_hake.py b/src/echodataflow/flows/flows_predict_hake.py index db750171..0ecda347 100644 --- a/src/echodataflow/flows/flows_predict_hake.py +++ b/src/echodataflow/flows/flows_predict_hake.py @@ -10,7 +10,6 @@ from prefect import flow, get_client, get_run_logger, runtime from prefect.states import Failed -from echodataflow.flows.flows_helper import cancel_if_deployment_already_running from echodataflow.utils.manifests import ( MVBS_COLUMNS_POSTPROCESSING, PREDICTION_COLUMNS_POSTPROCESSING, @@ -261,9 +260,6 @@ def flow_predict_hake_postprocessing( file_prediction_csv: str = "prediction_files.csv", ) -> None: """Predict all newly ready windows, combining aligned MVBS slices.""" - - if cancel_if_deployment_already_running(): - return logger = get_run_logger() file_MVBS_csv = Path(path_main) / file_MVBS_csv diff --git a/tests/deployment/conftest.py b/tests/deployment/conftest.py index e1beef3b..55869c83 100644 --- a/tests/deployment/conftest.py +++ b/tests/deployment/conftest.py @@ -1,5 +1,6 @@ import sys import types +from enum import Enum import pytest @@ -29,6 +30,16 @@ def __class_getitem__(cls, _item): return cls +class FakeConcurrencyLimitStrategy(str, Enum): + ENQUEUE = "ENQUEUE" + CANCEL_NEW = "CANCEL_NEW" + + +class FakeConcurrencyLimitConfig: + def __init__(self, **kwargs): + self.__dict__.update(kwargs) + + @pytest.fixture def install_prefect_stubs(monkeypatch): def _install(*, sink=None): @@ -36,6 +47,7 @@ def _install(*, sink=None): if sink is None: prefect_mod.deploy = lambda *args, **kwargs: None else: + def fake_deploy(*deployments, **kwargs): sink["deploy_call"] = { "deployments": list(deployments), @@ -57,12 +69,21 @@ def fake_deploy(*deployments, **kwargs): events_mod = types.ModuleType("prefect.events") events_mod.DeploymentEventTrigger = FakeTrigger + client_mod = types.ModuleType("prefect.client") + schemas_mod = types.ModuleType("prefect.client.schemas") + objects_mod = types.ModuleType("prefect.client.schemas.objects") + objects_mod.ConcurrencyLimitConfig = FakeConcurrencyLimitConfig + objects_mod.ConcurrencyLimitStrategy = FakeConcurrencyLimitStrategy + monkeypatch.setitem(sys.modules, "prefect", prefect_mod) monkeypatch.setitem(sys.modules, "prefect.deployments", deployments_mod) monkeypatch.setitem(sys.modules, "prefect.deployments.runner", runner_mod) monkeypatch.setitem(sys.modules, "prefect.flows", flows_mod) monkeypatch.setitem(sys.modules, "prefect.variables", variables_mod) monkeypatch.setitem(sys.modules, "prefect.events", events_mod) + monkeypatch.setitem(sys.modules, "prefect.client", client_mod) + monkeypatch.setitem(sys.modules, "prefect.client.schemas", schemas_mod) + monkeypatch.setitem(sys.modules, "prefect.client.schemas.objects", objects_mod) return { "FakeVariable": FakeVariable, @@ -71,4 +92,4 @@ def fake_deploy(*deployments, **kwargs): "FakePrefectFlowGeneric": FakePrefectFlowGeneric, } - return _install \ No newline at end of file + return _install diff --git a/tests/deployment/test_deploy_engine.py b/tests/deployment/test_deploy_engine.py index 1c1f02e7..30984b53 100644 --- a/tests/deployment/test_deploy_engine.py +++ b/tests/deployment/test_deploy_engine.py @@ -202,6 +202,7 @@ def from_source(self, **kwargs): entrypoint="echodataflow/flows/flows_acoustics.py:flow_raw2Sv_postprocessing", parameters={}, concurrency_group="postprocessing", + concurrency_limit={"limit": 1, "collision_strategy": "CANCEL_NEW"}, task_runner={"type": "dask", "cluster_kwargs": {"n_workers": 4}}, ) @@ -212,6 +213,9 @@ def from_source(self, **kwargs): ) assert calls["deployment"]["work_queue_name"] == "postprocessing" + native_limit = calls["deployment"]["concurrency_limit"] + assert native_limit.limit == 1 + assert native_limit.collision_strategy.value == "CANCEL_NEW" runtime_config = calls["deployment"]["job_variables"]["env"]["ECHODATAFLOW_TASK_RUNNER"] assert runtime_config == ('{"type": "dask", "cluster_kwargs": {"n_workers": 4}}') assert len(grouped) == 1 @@ -278,6 +282,10 @@ def test_validate_deploy_config_accepts_every_allowed_key(install_prefect_stubs) "interval": 10, "cron_offset": 3, "inject_time_offset": True, + "concurrency_limit": { + "limit": 1, + "collision_strategy": "CANCEL_NEW", + }, "work_pool_name": "special-pool", "task_runner": { "type": "dask", @@ -314,12 +322,18 @@ def test_validate_deploy_config_accepts_every_allowed_key(install_prefect_stubs) "flow", "interval", "cron_offset", + "concurrency_limit", "triggers", "inject_time_offset", "task_runner", "work_pool_name", } assert core.ALLOWED_CONCURRENCY_GROUP_KEYS == {"limit"} + assert core.ALLOWED_CONCURRENCY_LIMIT_KEYS == { + "limit", + "collision_strategy", + "grace_period_seconds", + } assert core.ALLOWED_TASK_RUNNER_KEYS == {"type", "cluster_kwargs"} assert core.ALLOWED_DASK_CLUSTER_KEYS == { "memory_limit", @@ -586,3 +600,40 @@ class _FakeFlow: deploy_cfg=deploy_cfg, resolved_flows=resolved_flows, ) + + +def test_build_deploy_specs_preserves_native_concurrency_limit(install_prefect_stubs): + install_prefect_stubs() + engine = importlib.import_module("echodataflow.deployment.deployment_engine") + + def _flow_fn(): + return None + + class _FakeFlow: + fn = staticmethod(_flow_fn) + + specs = engine.build_deploy_specs( + param_cfg={"flows": {"raw2Sv": {}}}, + deploy_cfg={ + "flows": { + "raw2Sv": { + "concurrency_limit": { + "limit": 1, + "collision_strategy": "CANCEL_NEW", + }, + } + } + }, + resolved_flows={ + "raw2Sv": { + "flow_obj": _FakeFlow(), + "entrypoint": "echodataflow/flows/flows_acoustics.py:flow_raw2Sv", + } + }, + ) + + assert specs[0].parameters == {} + assert specs[0].concurrency_limit == { + "limit": 1, + "collision_strategy": "CANCEL_NEW", + } diff --git a/tests/test_postprocessing_flow_validation.py b/tests/test_postprocessing_flow_validation.py index a480df0d..06a95136 100644 --- a/tests/test_postprocessing_flow_validation.py +++ b/tests/test_postprocessing_flow_validation.py @@ -3,7 +3,6 @@ from echodataflow.flows.flows_acoustics import ( flow_create_MVBS_postprocessing, - flow_raw2Sv, ) from echodataflow.flows.flows_predict_hake import flow_predict_hake_postprocessing @@ -16,12 +15,6 @@ def info(self, message): messages.append(message) monkeypatch.setattr(flows_acoustics, "get_run_logger", lambda: Logger()) - monkeypatch.setattr( - flows_acoustics, - "cancel_if_deployment_already_running", - lambda: False, - ) - flow_create_MVBS_postprocessing.fn(path_main=str(tmp_path)) assert messages == ["Sv ledger does not yet exist"] @@ -36,12 +29,6 @@ def info(self, message): messages.append(message) monkeypatch.setattr(flows_predict_hake, "get_run_logger", lambda: Logger()) - monkeypatch.setattr( - flows_predict_hake, - "cancel_if_deployment_already_running", - lambda: False, - ) - flow_predict_hake_postprocessing.fn( path_main=str(tmp_path), path_weight="unused.ckpt", @@ -49,40 +36,3 @@ def info(self, message): assert messages == ["MVBS ledger does not yet exist"] assert not (tmp_path / "prediction_files.csv").exists() - - -def test_mvbs_flow_exits_when_deployment_is_already_running(tmp_path, monkeypatch): - monkeypatch.setattr( - flows_acoustics, - "cancel_if_deployment_already_running", - lambda: True, - ) - - flow_create_MVBS_postprocessing.fn(path_main=str(tmp_path)) - - assert not (tmp_path / "MVBS_files.csv").exists() - - -def test_realtime_raw2sv_flow_exits_when_deployment_is_already_running(monkeypatch): - monkeypatch.setattr( - flows_acoustics, - "cancel_if_deployment_already_running", - lambda: True, - ) - - assert flow_raw2Sv.fn() is None - - -def test_prediction_flow_exits_when_deployment_is_already_running(tmp_path, monkeypatch): - monkeypatch.setattr( - flows_predict_hake, - "cancel_if_deployment_already_running", - lambda: True, - ) - - flow_predict_hake_postprocessing.fn( - path_main=str(tmp_path), - path_weight="unused.ckpt", - ) - - assert not (tmp_path / "prediction_files.csv").exists() From 7ae5243cc82f74403c05f0ce8af6fe87b4423c3f Mon Sep 17 00:00:00 2001 From: Wu-Jung Lee Date: Fri, 14 Aug 2026 22:07:46 -0700 Subject: [PATCH 4/8] add max_flow_run_attempts, change default runner to ThreadPoolTaskRunner --- src/echodataflow/deployment/task_runners.py | 7 +- src/echodataflow/flows/flows_acoustics.py | 56 ++++++++++++++-- src/echodataflow/flows/flows_predict_hake.py | 16 +++++ .../operations/operations_postprocessing.py | 57 +++++++++++++++- src/echodataflow/utils/manifests.py | 2 + tests/deployment/test_task_runners.py | 12 ++-- tests/test_operations_postprocessing.py | 66 +++++++++++++++++++ tests/test_postprocessing_flow_validation.py | 66 +++++++++++++++++++ 8 files changed, 265 insertions(+), 17 deletions(-) diff --git a/src/echodataflow/deployment/task_runners.py b/src/echodataflow/deployment/task_runners.py index 2b27674d..bd59e723 100644 --- a/src/echodataflow/deployment/task_runners.py +++ b/src/echodataflow/deployment/task_runners.py @@ -6,15 +6,16 @@ import os from prefect_dask import DaskTaskRunner +from prefect.task_runners import ThreadPoolTaskRunner from echodataflow.deployment.core import TASK_RUNNER_ENV_VAR -def dask_task_runner_from_environment() -> DaskTaskRunner: - """Build the flow's Dask runner when its entrypoint is loaded by a worker.""" +def task_runner_from_environment() -> DaskTaskRunner | ThreadPoolTaskRunner: + """Build the configured runner, or use Prefect's default thread pool.""" serialized = os.getenv(TASK_RUNNER_ENV_VAR) if serialized is None: - return DaskTaskRunner() + return ThreadPoolTaskRunner() config = json.loads(serialized) if config.get("type") != "dask": diff --git a/src/echodataflow/flows/flows_acoustics.py b/src/echodataflow/flows/flows_acoustics.py index 3a7e7bd8..9f34c1f9 100644 --- a/src/echodataflow/flows/flows_acoustics.py +++ b/src/echodataflow/flows/flows_acoustics.py @@ -13,7 +13,7 @@ from prefect.states import Failed from prefect import runtime -from echodataflow.deployment.task_runners import dask_task_runner_from_environment +from echodataflow.deployment.task_runners import task_runner_from_environment from echodataflow.utils.manifests import ( MVBS_COLUMNS_POSTPROCESSING, MVBS_COLUMNS_REALTIME, @@ -34,6 +34,8 @@ from echodataflow.operations.operations_postprocessing import ( build_MVBS_ledger, build_Sv_ledger, + failure_state, + propagate_blocked_status, plan_mvbs_slices, read_or_create_ledger, ) @@ -52,7 +54,7 @@ ep.utils.log.verbose() -@flow(log_prints=True, task_runner=dask_task_runner_from_environment()) +@flow(log_prints=True, task_runner=task_runner_from_environment()) def flow_raw2Sv( exclude_before: str | None = None, exclude_raw_file: list[str] = [], @@ -388,16 +390,18 @@ async def flow_create_MVBS( raise Exception(error_msg) -@flow(log_prints=True, task_runner=dask_task_runner_from_environment()) +@flow(log_prints=True, task_runner=task_runner_from_environment()) def flow_raw2Sv_postprocessing( path_raw_list: str, path_raw: str, path_main: str, start_time: str | None = None, end_time: str | None = None, + exclude_raw_file: list[str] = [], new_file_num_limit: int = -1, task_retries: int = 3, task_retry_delay_seconds: int = 30, + max_flow_run_attempts: int = 3, encode_mode: str = "power", waveform_mode: str = "CW", depth_offset: float = 9.5, @@ -430,7 +434,13 @@ def flow_raw2Sv_postprocessing( start_time=start_time, end_time=end_time, ) - selected = selected[selected["raw2Sv_status"] != "completed"] + + # Skip files explicitly excluded from conversion + if exclude_raw_file: + print(f"Exclude {exclude_raw_file} from processing") + selected = selected[~selected["raw_filename"].isin(exclude_raw_file)] + + selected = selected[selected["raw2Sv_status"].isin(["pending", "failed"])] if new_file_num_limit != -1: selected = selected.head(new_file_num_limit) if selected.empty: @@ -495,17 +505,26 @@ def flow_raw2Sv_postprocessing( logger.info("Completed %s", key) except Exception as exc: errors.append(exc) - df_Sv.loc[idx, ["raw2Sv_status", "error"]] = ["failed", str(exc)] + attempt_count, status = failure_state( + int(df_Sv.loc[idx, "attempt_count"]), max_flow_run_attempts + ) + df_Sv.loc[idx, ["raw2Sv_status", "attempt_count", "error"]] = [ + status, + attempt_count, + str(exc), + ] write_manifest(df_Sv, file_Sv_csv) logger.error("Failed to convert %s: %s", key, exc) if errors: raise RuntimeError(f"{len(errors)} raw-to-Sv conversions failed") -@flow(log_prints=True, task_runner=dask_task_runner_from_environment()) +@flow(log_prints=True, task_runner=task_runner_from_environment()) def flow_create_MVBS_postprocessing( path_main: str, slice_mins: int = 20, + new_file_num_limit: int = -1, + max_flow_run_attempts: int = 3, range_bin: str = "1m", ping_time_bin: str = "5s", file_Sv_csv: str = "Sv_files.csv", @@ -536,8 +555,24 @@ def flow_create_MVBS_postprocessing( builder=lambda: build_MVBS_ledger(df_Sv, slice_mins), ) + # Make terminal raw failures explicit in downstream planning + updated_MVBS = propagate_blocked_status( + df_Sv, + df_MVBS, + upstream_filename_column="raw_filename", + upstream_status_column="raw2Sv_status", + downstream_filenames_column="raw_filenames", + downstream_status_column="MVBS_status", + upstream_label="raw files", + ) + if not updated_MVBS.equals(df_MVBS): + df_MVBS = updated_MVBS + write_manifest(df_MVBS.sort_values("slice_start"), file_MVBS_csv) + # Get MVBS slices to be computed based on raw-to-Sv completions planned = plan_mvbs_slices(df_Sv, df_MVBS) + if new_file_num_limit != -1: + planned = planned[:new_file_num_limit] if not planned: logger.info("No newly ready MVBS slices") return @@ -599,7 +634,14 @@ def flow_create_MVBS_postprocessing( errors.append(exc) filename = f"MVBS_{item.start_time:%Y%m%dT%H%M%S}.zarr" idx = df_MVBS.index[df_MVBS["MVBS_filename"] == filename][0] - df_MVBS.loc[idx, ["MVBS_status", "error"]] = ["failed", str(exc)] + attempt_count, status = failure_state( + int(df_MVBS.loc[idx, "attempt_count"]), max_flow_run_attempts + ) + df_MVBS.loc[idx, ["MVBS_status", "attempt_count", "error"]] = [ + status, + attempt_count, + str(exc), + ] logger.error("MVBS slice %s failed: %s", item.start_time, exc) write_manifest(df_MVBS.sort_values("slice_start"), file_MVBS_csv) if errors: diff --git a/src/echodataflow/flows/flows_predict_hake.py b/src/echodataflow/flows/flows_predict_hake.py index 0ecda347..9f3348b0 100644 --- a/src/echodataflow/flows/flows_predict_hake.py +++ b/src/echodataflow/flows/flows_predict_hake.py @@ -20,6 +20,7 @@ ) from echodataflow.operations.operations_postprocessing import ( build_prediction_ledger, + propagate_blocked_status, plan_prediction_slices, read_or_create_ledger, ) @@ -287,6 +288,21 @@ def flow_predict_hake_postprocessing( prediction_slice_mins, ), ) + + # Make terminal MVBS failures explicit in downstream planning + updated_prediction = propagate_blocked_status( + df_MVBS, + df_prediction, + upstream_filename_column="MVBS_filename", + upstream_status_column="MVBS_status", + downstream_filenames_column="MVBS_filenames", + downstream_status_column="prediction_status", + upstream_label="MVBS slices", + ) + if not updated_prediction.equals(df_prediction): + df_prediction = updated_prediction + write_manifest(df_prediction.sort_values("slice_start"), file_prediction_csv) + # Assemble aligned prediction windows from completed MVBS slices planned = plan_prediction_slices( df_MVBS, diff --git a/src/echodataflow/operations/operations_postprocessing.py b/src/echodataflow/operations/operations_postprocessing.py index 4acec330..610cfec5 100644 --- a/src/echodataflow/operations/operations_postprocessing.py +++ b/src/echodataflow/operations/operations_postprocessing.py @@ -131,6 +131,7 @@ def build_Sv_ledger(raw_files: pd.DataFrame) -> pd.DataFrame: ledger["Sv_filename"] = pd.NA ledger["raw2Sv_status"] = "pending" + ledger["attempt_count"] = 0 ledger["first_ping_time"] = pd.NaT ledger["last_ping_time"] = pd.NaT ledger["error"] = "" @@ -172,12 +173,64 @@ def build_MVBS_ledger(ledger_Sv: pd.DataFrame, slice_mins: int = 20) -> pd.DataF "last_ping_time": pd.NaT, "is_partial": pd.NA, "MVBS_status": "pending", + "attempt_count": 0, "error": "", } ) return pd.DataFrame.from_records(records, columns=MVBS_COLUMNS_POSTPROCESSING) +def failure_state(attempt_count: int, max_flow_run_attempts: int) -> tuple[int, str]: + """Increment a failure count and return its retryable or terminal status.""" + if max_flow_run_attempts < 1: + raise ValueError("max_flow_run_attempts must be a positive integer") + + attempt_count += 1 + status = "always_failed" if attempt_count >= max_flow_run_attempts else "failed" + return attempt_count, status + + +def propagate_blocked_status( + upstream_ledger: pd.DataFrame, + downstream_ledger: pd.DataFrame, + *, + upstream_filename_column: str, + upstream_status_column: str, + downstream_filenames_column: str, + downstream_status_column: str, + upstream_label: str, +) -> pd.DataFrame: + """Mark downstream rows blocked by terminal upstream failures.""" + updated = downstream_ledger.copy() + if upstream_ledger.empty or downstream_ledger.empty: + return updated + + # Collect upstream files that are blocked or always_failed + terminal_inputs = set( + upstream_ledger.loc[ + upstream_ledger[upstream_status_column].isin(["blocked", "always_failed"]), + upstream_filename_column, + ].astype(str) + ) + if not terminal_inputs: + return updated + + # Only pending or retryable rows can transition to blocked + for idx, row in updated.iterrows(): + if row[downstream_status_column] not in {"pending", "failed"}: + continue + blocked_by = sorted( + terminal_inputs.intersection(json.loads(row[downstream_filenames_column])) + ) + # Record the terminal inputs that prevent downstream processing + if blocked_by: + updated.loc[idx, [downstream_status_column, "error"]] = [ + "blocked", + f"Required {upstream_label} cannot be processed: {blocked_by}", + ] + return updated + + def build_prediction_ledger( ledger_MVBS: pd.DataFrame, prediction_slice_mins: int = 40, @@ -259,7 +312,7 @@ def plan_mvbs_slices( slices: list[PlannedSlice] = [] for row in df_MVBS.itertuples(index=False): # Skip slices that have already been created successfully - if row.MVBS_status in {"completed", "no_data"}: + if row.MVBS_status not in {"pending", "failed"}: continue required_raw = json.loads(row.raw_filenames) @@ -302,7 +355,7 @@ def plan_prediction_slices( slices: list[PlannedSlice] = [] for row in df_prediction.itertuples(index=False): - if row.prediction_status == "completed": + if row.prediction_status not in {"pending", "failed"}: continue required_names = json.loads(row.MVBS_filenames) diff --git a/src/echodataflow/utils/manifests.py b/src/echodataflow/utils/manifests.py index 735c454c..35074d9e 100644 --- a/src/echodataflow/utils/manifests.py +++ b/src/echodataflow/utils/manifests.py @@ -27,6 +27,7 @@ "raw_filename", "Sv_filename", "raw2Sv_status", + "attempt_count", "first_ping_time", "last_ping_time", "error", @@ -40,6 +41,7 @@ "last_ping_time", "is_partial", "MVBS_status", + "attempt_count", "error", ] PREDICTION_COLUMNS_POSTPROCESSING = [ diff --git a/tests/deployment/test_task_runners.py b/tests/deployment/test_task_runners.py index 85d48f32..bd9a17b6 100644 --- a/tests/deployment/test_task_runners.py +++ b/tests/deployment/test_task_runners.py @@ -1,8 +1,10 @@ import json +from prefect.task_runners import ThreadPoolTaskRunner + from echodataflow.deployment.task_runners import ( TASK_RUNNER_ENV_VAR, - dask_task_runner_from_environment, + task_runner_from_environment, ) @@ -21,7 +23,7 @@ def test_dask_task_runner_uses_runtime_environment(monkeypatch): ), ) - runner = dask_task_runner_from_environment() + runner = task_runner_from_environment() assert runner.cluster_kwargs == { "n_workers": 4, @@ -30,9 +32,9 @@ def test_dask_task_runner_uses_runtime_environment(monkeypatch): } -def test_dask_task_runner_preserves_default_without_runtime_config(monkeypatch): +def test_task_runner_uses_prefect_default_without_runtime_config(monkeypatch): monkeypatch.delenv(TASK_RUNNER_ENV_VAR, raising=False) - runner = dask_task_runner_from_environment() + runner = task_runner_from_environment() - assert runner.cluster_kwargs == {} + assert isinstance(runner, ThreadPoolTaskRunner) diff --git a/tests/test_operations_postprocessing.py b/tests/test_operations_postprocessing.py index d30e75a8..d395cfd4 100644 --- a/tests/test_operations_postprocessing.py +++ b/tests/test_operations_postprocessing.py @@ -8,7 +8,9 @@ build_prediction_ledger, build_MVBS_ledger, build_Sv_ledger, + failure_state, generate_aligned_windows, + propagate_blocked_status, plan_mvbs_slices, plan_prediction_slices, select_contained_records, @@ -131,6 +133,38 @@ def test_ledgers_predeclare_raw_files_and_mvbs_slices(): '"IWCPS-D20250611-T001400.raw"]', '["IWCPS-D20250611-T001400.raw", "IWCPS-D20250611-T002100.raw"]', ] + assert sv["attempt_count"].tolist() == [0, 0, 0, 0] + assert mvbs["attempt_count"].tolist() == [0, 0] + + +def test_failure_state_becomes_terminal_at_configured_attempt(): + assert failure_state(0, 3) == (1, "failed") + assert failure_state(2, 3) == (3, "always_failed") + + with pytest.raises(ValueError, match="positive integer"): + failure_state(0, 0) + + +def test_terminal_raw_failure_blocks_dependent_mvbs_slice(): + df_Sv = _sv_ledger() + df_Sv.loc[1, "raw2Sv_status"] = "always_failed" + df_MVBS = build_MVBS_ledger(df_Sv, slice_mins=20) + + blocked = propagate_blocked_status( + df_Sv, + df_MVBS, + upstream_filename_column="raw_filename", + upstream_status_column="raw2Sv_status", + downstream_filenames_column="raw_filenames", + downstream_status_column="MVBS_status", + upstream_label="raw files", + ) + + assert blocked.loc[0, "MVBS_status"] == "blocked" + assert "IWCPS-D20250611-T000700.raw" in blocked.loc[0, "error"] + assert blocked.loc[1, "MVBS_status"] == "pending" + planned = plan_mvbs_slices(df_Sv, blocked) + assert [item.start_time for item in planned] == [pd.Timestamp("2025-06-11T00:20:00Z")] def test_sv_ledger_derives_timestamp_only_from_s3_path(): @@ -212,6 +246,38 @@ def test_prediction_planner_skips_completed_prediction_windows(): assert planned == [] +def test_terminal_mvbs_failure_blocks_dependent_prediction_window(): + starts = pd.date_range("2025-06-11T00:00:00Z", periods=4, freq="20min") + mvbs = pd.DataFrame( + { + "MVBS_filename": [f"slice-{index}.zarr" for index in range(4)], + "slice_start": starts, + "slice_end": starts + pd.Timedelta(minutes=20), + "first_ping_time": starts, + "last_ping_time": starts + pd.Timedelta(minutes=20), + "is_partial": [False] * 4, + "MVBS_status": ["blocked", "completed", "completed", "completed"], + } + ) + prediction = build_prediction_ledger(mvbs, 40) + + blocked = propagate_blocked_status( + mvbs, + prediction, + upstream_filename_column="MVBS_filename", + upstream_status_column="MVBS_status", + downstream_filenames_column="MVBS_filenames", + downstream_status_column="prediction_status", + upstream_label="MVBS slices", + ) + + assert blocked.loc[0, "prediction_status"] == "blocked" + assert "slice-0.zarr" in blocked.loc[0, "error"] + assert blocked.loc[1, "prediction_status"] == "pending" + planned = plan_prediction_slices(mvbs, blocked) + assert [item.start_time for item in planned] == [pd.Timestamp("2025-06-11T00:40:00Z")] + + def test_prediction_requires_both_mvbs_slices_by_default(): mvbs = pd.DataFrame( { diff --git a/tests/test_postprocessing_flow_validation.py b/tests/test_postprocessing_flow_validation.py index 06a95136..2c02cb4b 100644 --- a/tests/test_postprocessing_flow_validation.py +++ b/tests/test_postprocessing_flow_validation.py @@ -1,8 +1,11 @@ +import pandas as pd + import echodataflow.flows.flows_acoustics as flows_acoustics import echodataflow.flows.flows_predict_hake as flows_predict_hake from echodataflow.flows.flows_acoustics import ( flow_create_MVBS_postprocessing, + flow_raw2Sv_postprocessing, ) from echodataflow.flows.flows_predict_hake import flow_predict_hake_postprocessing @@ -36,3 +39,66 @@ def info(self, message): assert messages == ["MVBS ledger does not yet exist"] assert not (tmp_path / "prediction_files.csv").exists() + + +def test_raw2sv_postprocessing_skips_excluded_raw_files(tmp_path, monkeypatch, capsys): + raw_filename = "IWCPS-D20250619-T000204.raw" + ledger = pd.DataFrame( + { + "s3_path": [f"survey/{raw_filename}"], + "timestamp": [pd.Timestamp("2025-06-19T00:02:04Z")], + "raw_filename": [raw_filename], + "Sv_filename": [pd.NA], + "raw2Sv_status": ["pending"], + "first_ping_time": [pd.NaT], + "last_ping_time": [pd.NaT], + "error": [""], + } + ) + messages = [] + + class Logger: + def info(self, message): + messages.append(message) + + monkeypatch.setattr(flows_acoustics, "get_run_logger", lambda: Logger()) + monkeypatch.setattr( + flows_acoustics, + "read_or_create_ledger", + lambda **_: ledger.copy(), + ) + + flow_raw2Sv_postprocessing.fn( + path_raw_list="unused.csv", + path_raw=str(tmp_path), + path_main=str(tmp_path), + exclude_raw_file=[raw_filename], + ) + + assert messages == ["No raw files require processing"] + assert f"Exclude ['{raw_filename}'] from processing" in capsys.readouterr().out + + +def test_mvbs_postprocessing_applies_new_file_num_limit(tmp_path, monkeypatch): + (tmp_path / "Sv_files.csv").touch() + messages = [] + + class Logger: + def info(self, message): + messages.append(message) + + monkeypatch.setattr(flows_acoustics, "get_run_logger", lambda: Logger()) + monkeypatch.setattr(flows_acoustics, "read_manifest", lambda **_: pd.DataFrame()) + monkeypatch.setattr( + flows_acoustics, + "read_or_create_ledger", + lambda **_: pd.DataFrame(), + ) + monkeypatch.setattr(flows_acoustics, "plan_mvbs_slices", lambda *_: [object()]) + + flow_create_MVBS_postprocessing.fn( + path_main=str(tmp_path), + new_file_num_limit=0, + ) + + assert messages == ["No newly ready MVBS slices"] From f0d110001085987219302dd6b776c0e1f7d00adb Mon Sep 17 00:00:00 2001 From: Wu-Jung Lee Date: Sat, 15 Aug 2026 08:34:18 -0700 Subject: [PATCH 5/8] add new_file_num_limit also to predict hake flow --- src/echodataflow/flows/flows_predict_hake.py | 3 ++ tests/test_postprocessing_flow_validation.py | 30 ++++++++++++++++++++ 2 files changed, 33 insertions(+) diff --git a/src/echodataflow/flows/flows_predict_hake.py b/src/echodataflow/flows/flows_predict_hake.py index 9f3348b0..20a029a2 100644 --- a/src/echodataflow/flows/flows_predict_hake.py +++ b/src/echodataflow/flows/flows_predict_hake.py @@ -252,6 +252,7 @@ def flow_predict_hake_postprocessing( path_main: str, path_weight: str, prediction_slice_mins: int = 40, + new_file_num_limit: int = -1, temperature: float = 0.5, softmax_threshold: float = 0.5, max_depth: float = 590.0, @@ -308,6 +309,8 @@ def flow_predict_hake_postprocessing( df_MVBS, df_prediction, ) + if new_file_num_limit != -1: + planned = planned[:new_file_num_limit] if not planned: logger.info("No newly ready prediction windows") return diff --git a/tests/test_postprocessing_flow_validation.py b/tests/test_postprocessing_flow_validation.py index 2c02cb4b..ecf76a9d 100644 --- a/tests/test_postprocessing_flow_validation.py +++ b/tests/test_postprocessing_flow_validation.py @@ -102,3 +102,33 @@ def info(self, message): ) assert messages == ["No newly ready MVBS slices"] + + +def test_prediction_postprocessing_applies_new_file_num_limit(tmp_path, monkeypatch): + (tmp_path / "MVBS_files.csv").touch() + messages = [] + + class Logger: + def info(self, message): + messages.append(message) + + monkeypatch.setattr(flows_predict_hake, "get_run_logger", lambda: Logger()) + monkeypatch.setattr(flows_predict_hake, "read_manifest", lambda *_, **__: pd.DataFrame()) + monkeypatch.setattr( + flows_predict_hake, + "read_or_create_ledger", + lambda **_: pd.DataFrame(), + ) + monkeypatch.setattr( + flows_predict_hake, + "plan_prediction_slices", + lambda *_: [object()], + ) + + flow_predict_hake_postprocessing.fn( + path_main=str(tmp_path), + path_weight="unused.ckpt", + new_file_num_limit=0, + ) + + assert messages == ["No newly ready prediction windows"] From 1cf8e3d0bbf66ab3d24a9473b46c6ff31207958c Mon Sep 17 00:00:00 2001 From: Wu-Jung Lee Date: Sat, 15 Aug 2026 14:52:19 -0700 Subject: [PATCH 6/8] add changes to remove Sv files for which all associated MVBS files are processed --- src/echodataflow/flows/flows_acoustics.py | 90 ++++++++++++++----- .../operations/operations_postprocessing.py | 29 ++++++ src/echodataflow/utils/manifests.py | 3 + tests/test_operations_postprocessing.py | 22 +++++ tests/test_postprocessing_flow_validation.py | 35 +++++++- 5 files changed, 157 insertions(+), 22 deletions(-) diff --git a/src/echodataflow/flows/flows_acoustics.py b/src/echodataflow/flows/flows_acoustics.py index 9f34c1f9..3ee74470 100644 --- a/src/echodataflow/flows/flows_acoustics.py +++ b/src/echodataflow/flows/flows_acoustics.py @@ -3,6 +3,7 @@ import asyncio import datetime from pathlib import Path +import shutil import pandas as pd @@ -37,6 +38,7 @@ failure_state, propagate_blocked_status, plan_mvbs_slices, + plan_Sv_cleanup, read_or_create_ledger, ) from echodataflow.tasks.tasks_acoustics import ( @@ -424,7 +426,7 @@ def flow_raw2Sv_postprocessing( df_Sv = read_or_create_ledger( ledger_path=file_Sv_csv, columns=SV_COLUMNS_POSTPROCESSING, - date_columns=["timestamp", "first_ping_time", "last_ping_time"], + date_columns=["timestamp", "first_ping_time", "last_ping_time", "Sv_deleted_at"], builder=lambda: build_Sv_ledger(pd.read_csv(path_raw_list)), ) selected = filter_time_range( @@ -492,6 +494,9 @@ def flow_raw2Sv_postprocessing( "first_ping_time", "last_ping_time", "error", + "Sv_cleanup_status", + "Sv_deleted_at", + "Sv_cleanup_error", ], ] = [ result.filename_raw, @@ -500,6 +505,9 @@ def flow_raw2Sv_postprocessing( result.first_ping_time, result.last_ping_time, "", + "pending", + pd.NaT, + "", ] write_manifest(df_Sv, file_Sv_csv) logger.info("Completed %s", key) @@ -529,6 +537,7 @@ def flow_create_MVBS_postprocessing( ping_time_bin: str = "5s", file_Sv_csv: str = "Sv_files.csv", file_MVBS_csv: str = "MVBS_files.csv", + remove_completed_Sv_files: bool = True, ) -> None: """Create preplanned MVBS slices after all required raw conversions finish.""" @@ -545,7 +554,7 @@ def flow_create_MVBS_postprocessing( df_Sv = read_manifest( path=file_Sv_csv, columns=SV_COLUMNS_POSTPROCESSING, - date_columns=["timestamp", "first_ping_time", "last_ping_time"], + date_columns=["timestamp", "first_ping_time", "last_ping_time", "Sv_deleted_at"], ) # Preplan every MVBS row once from the complete raw-file timeline df_MVBS = read_or_create_ledger( @@ -575,28 +584,28 @@ def flow_create_MVBS_postprocessing( planned = planned[:new_file_num_limit] if not planned: logger.info("No newly ready MVBS slices") - return - settings = CreateMVBSSettings( - sv_directory=str(Path(path_main) / "Sv"), - output_directory=str(path_MVBS), - range_bin=range_bin, - ping_time_bin=ping_time_bin, - ) - # Ready slices are independent and can be computed in parallel futures = {} - for item in planned: - filename = f"MVBS_{item.start_time:%Y%m%dT%H%M%S}.zarr" - future = task_create_MVBS.with_options(task_run_name=filename).submit( - CreateMVBSWorkItem( - start_time=item.start_time, - end_time=item.end_time, - sv_filenames=item.filenames, - mvbs_filename=filename, - ), - settings, + if planned: + settings = CreateMVBSSettings( + sv_directory=str(Path(path_main) / "Sv"), + output_directory=str(path_MVBS), + range_bin=range_bin, + ping_time_bin=ping_time_bin, ) - futures[future] = item + # Ready slices are independent and can be computed in parallel + for item in planned: + filename = f"MVBS_{item.start_time:%Y%m%dT%H%M%S}.zarr" + future = task_create_MVBS.with_options(task_run_name=filename).submit( + CreateMVBSWorkItem( + start_time=item.start_time, + end_time=item.end_time, + sv_filenames=item.filenames, + mvbs_filename=filename, + ), + settings, + ) + futures[future] = item # Collect task results in memory so this flow remains the sole manifest writer errors = [] @@ -644,5 +653,44 @@ def flow_create_MVBS_postprocessing( ] logger.error("MVBS slice %s failed: %s", item.start_time, exc) write_manifest(df_MVBS.sort_values("slice_start"), file_MVBS_csv) + + if remove_completed_Sv_files: + # Remove only Sv inputs whose dependent MVBS slices are all successful + path_Sv = (Path(path_main) / "Sv").resolve() + df_Sv["Sv_cleanup_error"] = df_Sv["Sv_cleanup_error"].fillna("").astype("string") + for Sv_filename in plan_Sv_cleanup(df_Sv, df_MVBS): + # Resolve the stable filename back to its single ledger row. + matches = df_Sv.index[df_Sv["Sv_filename"] == Sv_filename] + if len(matches) != 1: + raise ValueError( + f"Expected exactly one Sv ledger row for {Sv_filename}, found {len(matches)}" + ) + idx = matches[0] + file_Sv = (path_Sv / Sv_filename).resolve() + try: + # Guard recursive deletion against absolute, parent, or symlink escapes + if file_Sv.parent != path_Sv: + raise ValueError( + f"Sv filename resolves outside the Sv directory: {Sv_filename}" + ) + if file_Sv.is_dir(): + shutil.rmtree(file_Sv) + elif file_Sv.is_file(): + file_Sv.unlink() + else: + raise FileNotFoundError(f"Sv file does not exist: {file_Sv}") + df_Sv.loc[idx, "Sv_cleanup_status"] = "deleted" + df_Sv.loc[idx, "Sv_deleted_at"] = pd.Timestamp.now(tz="UTC") + df_Sv.loc[idx, "Sv_cleanup_error"] = "" + logger.info("Removed Sv file %s", file_Sv) + except Exception as exc: + df_Sv.loc[idx, ["Sv_cleanup_status", "Sv_cleanup_error"]] = [ + "failed", + str(exc), + ] + logger.error("Failed to remove Sv file %s: %s", file_Sv, exc) + # Persist each outcome so an interrupted run can safely resume cleanup + write_manifest(df_Sv.sort_values("timestamp"), file_Sv_csv) + if errors: raise RuntimeError(f"{len(errors)} MVBS slices failed") diff --git a/src/echodataflow/operations/operations_postprocessing.py b/src/echodataflow/operations/operations_postprocessing.py index 610cfec5..47514929 100644 --- a/src/echodataflow/operations/operations_postprocessing.py +++ b/src/echodataflow/operations/operations_postprocessing.py @@ -135,6 +135,9 @@ def build_Sv_ledger(raw_files: pd.DataFrame) -> pd.DataFrame: ledger["first_ping_time"] = pd.NaT ledger["last_ping_time"] = pd.NaT ledger["error"] = "" + ledger["Sv_cleanup_status"] = "pending" + ledger["Sv_deleted_at"] = pd.NaT + ledger["Sv_cleanup_error"] = "" return ledger.sort_values("timestamp").reset_index(drop=True) @@ -338,6 +341,32 @@ def plan_mvbs_slices( return slices +def plan_Sv_cleanup( + ledger_Sv: pd.DataFrame, + ledger_MVBS: pd.DataFrame, +) -> list[str]: + """Return Sv filenames safe to delete after all dependent MVBS work succeeds.""" + if ledger_Sv.empty or ledger_MVBS.empty: + return [] + + dependencies: dict[str, list[str]] = {} + for row in ledger_MVBS.itertuples(index=False): + for raw_filename in json.loads(row.raw_filenames): + dependencies.setdefault(raw_filename, []).append(row.MVBS_status) + + successful_statuses = {"completed", "no_data"} + cleanup_filenames = [] + for _, row in ledger_Sv.iterrows(): + if row["Sv_cleanup_status"] not in {"pending", "failed"}: + continue + if pd.isna(row["Sv_filename"]): + continue + statuses = dependencies.get(str(row["raw_filename"]), []) + if statuses and all(status in successful_statuses for status in statuses): + cleanup_filenames.append(str(row["Sv_filename"])) + return cleanup_filenames + + def plan_prediction_slices( ledger_MVBS: pd.DataFrame, ledger_prediction: pd.DataFrame, diff --git a/src/echodataflow/utils/manifests.py b/src/echodataflow/utils/manifests.py index 35074d9e..931ca40c 100644 --- a/src/echodataflow/utils/manifests.py +++ b/src/echodataflow/utils/manifests.py @@ -31,6 +31,9 @@ "first_ping_time", "last_ping_time", "error", + "Sv_cleanup_status", + "Sv_deleted_at", + "Sv_cleanup_error", ] MVBS_COLUMNS_POSTPROCESSING = [ "MVBS_filename", diff --git a/tests/test_operations_postprocessing.py b/tests/test_operations_postprocessing.py index d395cfd4..8aa90446 100644 --- a/tests/test_operations_postprocessing.py +++ b/tests/test_operations_postprocessing.py @@ -12,6 +12,7 @@ generate_aligned_windows, propagate_blocked_status, plan_mvbs_slices, + plan_Sv_cleanup, plan_prediction_slices, select_contained_records, select_overlapping_records, @@ -197,6 +198,27 @@ def test_pending_raw_input_keeps_its_mvbs_slice_closed(): assert plan_mvbs_slices(sv, build_MVBS_ledger(sv, slice_mins=20)) == [] +def test_sv_cleanup_waits_for_every_dependent_mvbs_slice(): + sv = _sv_ledger() + mvbs = build_MVBS_ledger(sv, slice_mins=20) + mvbs["MVBS_status"] = ["completed", "pending"] + + assert plan_Sv_cleanup(sv, mvbs) == ["a.zarr", "b.zarr"] + + mvbs.loc[1, "MVBS_status"] = "no_data" + + assert plan_Sv_cleanup(sv, mvbs) == ["a.zarr", "b.zarr", "c.zarr", "d.zarr"] + + +def test_sv_cleanup_skips_files_already_recorded_as_deleted(): + sv = _sv_ledger() + sv.loc[0, "Sv_cleanup_status"] = "deleted" + mvbs = build_MVBS_ledger(sv, slice_mins=20) + mvbs["MVBS_status"] = "completed" + + assert plan_Sv_cleanup(sv, mvbs) == ["b.zarr", "c.zarr", "d.zarr"] + + def test_empty_sv_ledger_builds_header_only_mvbs_ledger(): mvbs = build_MVBS_ledger(pd.DataFrame(), slice_mins=20) diff --git a/tests/test_postprocessing_flow_validation.py b/tests/test_postprocessing_flow_validation.py index ecf76a9d..90cb6ae9 100644 --- a/tests/test_postprocessing_flow_validation.py +++ b/tests/test_postprocessing_flow_validation.py @@ -3,6 +3,9 @@ import echodataflow.flows.flows_acoustics as flows_acoustics import echodataflow.flows.flows_predict_hake as flows_predict_hake +from echodataflow.operations.operations_postprocessing import build_MVBS_ledger, build_Sv_ledger +from echodataflow.utils.manifests import write_manifest + from echodataflow.flows.flows_acoustics import ( flow_create_MVBS_postprocessing, flow_raw2Sv_postprocessing, @@ -92,13 +95,14 @@ def info(self, message): monkeypatch.setattr( flows_acoustics, "read_or_create_ledger", - lambda **_: pd.DataFrame(), + lambda **_: pd.DataFrame(columns=["slice_start"]), ) monkeypatch.setattr(flows_acoustics, "plan_mvbs_slices", lambda *_: [object()]) flow_create_MVBS_postprocessing.fn( path_main=str(tmp_path), new_file_num_limit=0, + remove_completed_Sv_files=False, ) assert messages == ["No newly ready MVBS slices"] @@ -132,3 +136,32 @@ def info(self, message): ) assert messages == ["No newly ready prediction windows"] + + +def test_mvbs_postprocessing_removes_sv_after_all_dependencies_complete(tmp_path, monkeypatch): + sv = build_Sv_ledger(pd.DataFrame({"s3_path": ["survey/IWCPS-D20250611-T000100.raw"]})) + sv.loc[0, ["Sv_filename", "raw2Sv_status"]] = ["input.zarr", "completed"] + sv["first_ping_time"] = pd.to_datetime(["2025-06-11T00:01:00Z"], utc=True) + sv["last_ping_time"] = pd.to_datetime(["2025-06-11T00:10:00Z"], utc=True) + mvbs = build_MVBS_ledger(sv) + mvbs.loc[0, "MVBS_status"] = "completed" + write_manifest(sv, tmp_path / "Sv_files.csv") + write_manifest(mvbs, tmp_path / "MVBS_files.csv") + Sv_path = tmp_path / "Sv" / "input.zarr" + Sv_path.mkdir(parents=True) + + class Logger: + def info(self, *_args): + pass + + def error(self, *_args): + pass + + monkeypatch.setattr(flows_acoustics, "get_run_logger", lambda: Logger()) + + flow_create_MVBS_postprocessing.fn(path_main=str(tmp_path)) + + cleaned = pd.read_csv(tmp_path / "Sv_files.csv", index_col=0) + assert not Sv_path.exists() + assert cleaned.loc[0, "Sv_cleanup_status"] == "deleted" + assert pd.notna(cleaned.loc[0, "Sv_deleted_at"]) From 865bbc4957b0655f20b651b5d56d66525c14ddbf Mon Sep 17 00:00:00 2001 From: Wu-Jung Lee Date: Sat, 15 Aug 2026 21:21:49 -0700 Subject: [PATCH 7/8] assign no_data to the long gap on MVBS slices when building ledger --- src/echodataflow/flows/flows_acoustics.py | 7 ++- .../operations/operations_postprocessing.py | 24 ++++++++-- tests/test_operations_postprocessing.py | 47 +++++++++++++++++++ 3 files changed, 74 insertions(+), 4 deletions(-) diff --git a/src/echodataflow/flows/flows_acoustics.py b/src/echodataflow/flows/flows_acoustics.py index 3ee74470..b288c641 100644 --- a/src/echodataflow/flows/flows_acoustics.py +++ b/src/echodataflow/flows/flows_acoustics.py @@ -531,6 +531,7 @@ def flow_raw2Sv_postprocessing( def flow_create_MVBS_postprocessing( path_main: str, slice_mins: int = 20, + no_data_gap_hours: float = 3.0, new_file_num_limit: int = -1, max_flow_run_attempts: int = 3, range_bin: str = "1m", @@ -561,7 +562,11 @@ def flow_create_MVBS_postprocessing( ledger_path=file_MVBS_csv, columns=MVBS_COLUMNS_POSTPROCESSING, date_columns=["slice_start", "slice_end", "first_ping_time", "last_ping_time"], - builder=lambda: build_MVBS_ledger(df_Sv, slice_mins), + builder=lambda: build_MVBS_ledger( + ledger_Sv=df_Sv, + slice_mins=slice_mins, + no_data_gap_hours=no_data_gap_hours, + ), ) # Make terminal raw failures explicit in downstream planning diff --git a/src/echodataflow/operations/operations_postprocessing.py b/src/echodataflow/operations/operations_postprocessing.py index 47514929..a4c76e77 100644 --- a/src/echodataflow/operations/operations_postprocessing.py +++ b/src/echodataflow/operations/operations_postprocessing.py @@ -141,8 +141,14 @@ def build_Sv_ledger(raw_files: pd.DataFrame) -> pd.DataFrame: return ledger.sort_values("timestamp").reset_index(drop=True) -def build_MVBS_ledger(ledger_Sv: pd.DataFrame, slice_mins: int = 20) -> pd.DataFrame: +def build_MVBS_ledger( + ledger_Sv: pd.DataFrame, + slice_mins: int = 20, + no_data_gap_hours: float = 3.0, +) -> pd.DataFrame: """Preplan every MVBS slice and its required raw files.""" + if no_data_gap_hours <= 0: + raise ValueError("no_data_gap_hours must be greater than zero") if ledger_Sv.empty: return pd.DataFrame(columns=MVBS_COLUMNS_POSTPROCESSING) @@ -157,6 +163,7 @@ def build_MVBS_ledger(ledger_Sv: pd.DataFrame, slice_mins: int = 20) -> pd.DataF ) records = [] + no_data_gap = pd.Timedelta(hours=no_data_gap_hours) for window in windows: # Record the raw files required by this slice, including its predecessor required = filter_time_range( @@ -166,6 +173,13 @@ def build_MVBS_ledger(ledger_Sv: pd.DataFrame, slice_mins: int = 20) -> pd.DataF start_time=window.start_time, end_time=window.end_time, ) + + # Set the slice status to no_data if there is a long gap + # between the last Sv timestamp and the slice start + latest_Sv_timestamp = required["timestamp"].max() + is_long_gap = ( + pd.notna(latest_Sv_timestamp) and window.start_time - latest_Sv_timestamp > no_data_gap + ) records.append( { "MVBS_filename": f"MVBS_{window.start_time:%Y%m%dT%H%M%S}.zarr", @@ -175,9 +189,13 @@ def build_MVBS_ledger(ledger_Sv: pd.DataFrame, slice_mins: int = 20) -> pd.DataF "first_ping_time": pd.NaT, "last_ping_time": pd.NaT, "is_partial": pd.NA, - "MVBS_status": "pending", + "MVBS_status": "no_data" if is_long_gap else "pending", "attempt_count": 0, - "error": "", + "error": ( + f"No Sv file timestamp within {no_data_gap_hours} hours before the slice start" + if is_long_gap + else "" + ), } ) return pd.DataFrame.from_records(records, columns=MVBS_COLUMNS_POSTPROCESSING) diff --git a/tests/test_operations_postprocessing.py b/tests/test_operations_postprocessing.py index 8aa90446..fbd32e57 100644 --- a/tests/test_operations_postprocessing.py +++ b/tests/test_operations_postprocessing.py @@ -138,6 +138,53 @@ def test_ledgers_predeclare_raw_files_and_mvbs_slices(): assert mvbs["attempt_count"].tolist() == [0, 0] +def test_mvbs_ledger_marks_slices_more_than_three_hours_after_sv_as_no_data(): + sv = build_Sv_ledger( + pd.DataFrame( + { + "s3_path": [ + "survey/IWCPS-D20250611-T000000.raw", + "survey/IWCPS-D20250611-T040100.raw", + ] + } + ) + ) + + mvbs = build_MVBS_ledger(sv, slice_mins=20) + statuses = mvbs.set_index("slice_start")["MVBS_status"] + + # Exactly three hours remains eligible, while later empty gap slices do not. + assert statuses[pd.Timestamp("2025-06-11T03:00:00Z")] == "pending" + assert statuses[pd.Timestamp("2025-06-11T03:20:00Z")] == "no_data" + assert statuses[pd.Timestamp("2025-06-11T03:40:00Z")] == "no_data" + # The slice containing the first raw file after the gap remains eligible. + assert statuses[pd.Timestamp("2025-06-11T04:00:00Z")] == "pending" + + +def test_mvbs_ledger_accepts_custom_no_data_gap_hours(): + sv = build_Sv_ledger( + pd.DataFrame( + { + "s3_path": [ + "survey/IWCPS-D20250611-T000000.raw", + "survey/IWCPS-D20250611-T020100.raw", + ] + } + ) + ) + + mvbs = build_MVBS_ledger(sv, slice_mins=20, no_data_gap_hours=1.0) + statuses = mvbs.set_index("slice_start")["MVBS_status"] + + assert statuses[pd.Timestamp("2025-06-11T01:00:00Z")] == "pending" + assert statuses[pd.Timestamp("2025-06-11T01:20:00Z")] == "no_data" + + +def test_mvbs_ledger_rejects_nonpositive_no_data_gap_hours(): + with pytest.raises(ValueError, match="greater than zero"): + build_MVBS_ledger(pd.DataFrame(), no_data_gap_hours=0) + + def test_failure_state_becomes_terminal_at_configured_attempt(): assert failure_state(0, 3) == (1, "failed") assert failure_state(2, 3) == (3, "always_failed") From 7c847deda2b4094891c72ad29e5efe55e1e7c8e2 Mon Sep 17 00:00:00 2001 From: Wu-Jung Lee Date: Sun, 16 Aug 2026 09:15:29 -0700 Subject: [PATCH 8/8] update predict flow with changes in raw2Sv and create MVBS, refactor to create propagate_status function --- src/echodataflow/flows/flows_acoustics.py | 4 +- src/echodataflow/flows/flows_predict_hake.py | 15 +++- .../operations/operations_postprocessing.py | 54 ++++++++---- src/echodataflow/utils/manifests.py | 1 + tests/test_operations_postprocessing.py | 88 ++++++++++++++++++- 5 files changed, 136 insertions(+), 26 deletions(-) diff --git a/src/echodataflow/flows/flows_acoustics.py b/src/echodataflow/flows/flows_acoustics.py index b288c641..295552e4 100644 --- a/src/echodataflow/flows/flows_acoustics.py +++ b/src/echodataflow/flows/flows_acoustics.py @@ -36,7 +36,7 @@ build_MVBS_ledger, build_Sv_ledger, failure_state, - propagate_blocked_status, + propagate_status, plan_mvbs_slices, plan_Sv_cleanup, read_or_create_ledger, @@ -570,7 +570,7 @@ def flow_create_MVBS_postprocessing( ) # Make terminal raw failures explicit in downstream planning - updated_MVBS = propagate_blocked_status( + updated_MVBS = propagate_status( df_Sv, df_MVBS, upstream_filename_column="raw_filename", diff --git a/src/echodataflow/flows/flows_predict_hake.py b/src/echodataflow/flows/flows_predict_hake.py index 20a029a2..b8622fb1 100644 --- a/src/echodataflow/flows/flows_predict_hake.py +++ b/src/echodataflow/flows/flows_predict_hake.py @@ -20,8 +20,9 @@ ) from echodataflow.operations.operations_postprocessing import ( build_prediction_ledger, - propagate_blocked_status, + failure_state, plan_prediction_slices, + propagate_status, read_or_create_ledger, ) from echodataflow.operations.operations_predict_hake import ( @@ -253,6 +254,7 @@ def flow_predict_hake_postprocessing( path_weight: str, prediction_slice_mins: int = 40, new_file_num_limit: int = -1, + max_flow_run_attempts: int = 3, temperature: float = 0.5, softmax_threshold: float = 0.5, max_depth: float = 590.0, @@ -291,7 +293,7 @@ def flow_predict_hake_postprocessing( ) # Make terminal MVBS failures explicit in downstream planning - updated_prediction = propagate_blocked_status( + updated_prediction = propagate_status( df_MVBS, df_prediction, upstream_filename_column="MVBS_filename", @@ -381,7 +383,14 @@ def flow_predict_hake_postprocessing( except Exception as exc: errors.append(exc) idx = df_prediction.index[df_prediction["prediction_filename_postfix"] == postfix][0] - df_prediction.loc[idx, ["prediction_status", "error"]] = ["failed", str(exc)] + attempt_count, status = failure_state( + int(df_prediction.loc[idx, "attempt_count"]), max_flow_run_attempts + ) + df_prediction.loc[idx, ["prediction_status", "attempt_count", "error"]] = [ + status, + attempt_count, + str(exc), + ] write_manifest(df_prediction.sort_values("slice_start"), file_prediction_csv) logger.error("Prediction window %s failed: %s", item.start_time, exc) if errors: diff --git a/src/echodataflow/operations/operations_postprocessing.py b/src/echodataflow/operations/operations_postprocessing.py index a4c76e77..9b243a98 100644 --- a/src/echodataflow/operations/operations_postprocessing.py +++ b/src/echodataflow/operations/operations_postprocessing.py @@ -174,7 +174,7 @@ def build_MVBS_ledger( end_time=window.end_time, ) - # Set the slice status to no_data if there is a long gap + # Set the slice status to no_data if there is a long gap # between the last Sv timestamp and the slice start latest_Sv_timestamp = required["timestamp"].max() is_long_gap = ( @@ -211,7 +211,7 @@ def failure_state(attempt_count: int, max_flow_run_attempts: int) -> tuple[int, return attempt_count, status -def propagate_blocked_status( +def propagate_status( upstream_ledger: pd.DataFrame, downstream_ledger: pd.DataFrame, *, @@ -221,34 +221,44 @@ def propagate_blocked_status( downstream_status_column: str, upstream_label: str, ) -> pd.DataFrame: - """Mark downstream rows blocked by terminal upstream failures.""" + """Propagate terminal blocked and all-no-data upstream states downstream.""" updated = downstream_ledger.copy() if upstream_ledger.empty or downstream_ledger.empty: return updated - # Collect upstream files that are blocked or always_failed - terminal_inputs = set( + # A terminal failure in any dependency takes precedence over no-data. + blocked_inputs = set( upstream_ledger.loc[ upstream_ledger[upstream_status_column].isin(["blocked", "always_failed"]), upstream_filename_column, ].astype(str) ) - if not terminal_inputs: - return updated - - # Only pending or retryable rows can transition to blocked for idx, row in updated.iterrows(): if row[downstream_status_column] not in {"pending", "failed"}: continue - blocked_by = sorted( - terminal_inputs.intersection(json.loads(row[downstream_filenames_column])) - ) - # Record the terminal inputs that prevent downstream processing + required_names = json.loads(row[downstream_filenames_column]) + blocked_by = sorted(blocked_inputs.intersection(required_names)) if blocked_by: updated.loc[idx, [downstream_status_column, "error"]] = [ "blocked", f"Required {upstream_label} cannot be processed: {blocked_by}", ] + continue + + required_upstream = upstream_ledger[ + upstream_ledger[upstream_filename_column].isin(required_names) + ] + # All dependencies must be known before propagating no-data. + if len(required_upstream) != len(required_names): + continue + if ( + not required_upstream.empty + and required_upstream[upstream_status_column].eq("no_data").all() + ): + updated.loc[idx, [downstream_status_column, "error"]] = [ + "no_data", + f"All required {upstream_label} have no data", + ] return updated @@ -279,6 +289,7 @@ def build_prediction_ledger( end_time=window.end_time, include_exact_start_time=False, ) + is_no_data = not required.empty and required["MVBS_status"].eq("no_data").all() records.append( { "prediction_filename_postfix": window.start_time.strftime("%Y%m%dT%H%M%S"), @@ -291,8 +302,9 @@ def build_prediction_ledger( "evr_filename": pd.NA, "first_ping_time": pd.NaT, "last_ping_time": pd.NaT, - "prediction_status": "pending", - "error": "", + "prediction_status": "no_data" if is_no_data else "pending", + "attempt_count": 0, + "error": "No MVBS data in the prediction window" if is_no_data else "", } ) return pd.DataFrame.from_records(records, columns=PREDICTION_COLUMNS_POSTPROCESSING) @@ -410,13 +422,19 @@ def plan_prediction_slices( if len(required_MVBS) != len(required_names): raise ValueError(f"Prediction ledger references unknown MVBS files: {required_names}") - if not required_MVBS["MVBS_status"].eq("completed").all(): + # No-data slices are terminal and do not need to be passed to prediction. + if not required_MVBS["MVBS_status"].isin(["completed", "no_data"]).all(): + continue + + completed_MVBS = required_MVBS[required_MVBS["MVBS_status"] == "completed"] + if completed_MVBS.empty: continue partial = ( - required_MVBS["is_partial"] + completed_MVBS["is_partial"] .map(lambda value: str(value).strip().lower() in {"true", "1", "yes"}) .any() + or required_MVBS["MVBS_status"].eq("no_data").any() ) slices.append( @@ -424,7 +442,7 @@ def plan_prediction_slices( start_time=row.slice_start, end_time=row.slice_end, filenames=tuple( - required_MVBS.sort_values("slice_start")["MVBS_filename"].astype(str) + completed_MVBS.sort_values("slice_start")["MVBS_filename"].astype(str) ), is_partial=bool(partial), ) diff --git a/src/echodataflow/utils/manifests.py b/src/echodataflow/utils/manifests.py index 931ca40c..128d672e 100644 --- a/src/echodataflow/utils/manifests.py +++ b/src/echodataflow/utils/manifests.py @@ -59,6 +59,7 @@ "first_ping_time", "last_ping_time", "prediction_status", + "attempt_count", "error", ] diff --git a/tests/test_operations_postprocessing.py b/tests/test_operations_postprocessing.py index fbd32e57..7f36750a 100644 --- a/tests/test_operations_postprocessing.py +++ b/tests/test_operations_postprocessing.py @@ -10,7 +10,7 @@ build_Sv_ledger, failure_state, generate_aligned_windows, - propagate_blocked_status, + propagate_status, plan_mvbs_slices, plan_Sv_cleanup, plan_prediction_slices, @@ -198,7 +198,7 @@ def test_terminal_raw_failure_blocks_dependent_mvbs_slice(): df_Sv.loc[1, "raw2Sv_status"] = "always_failed" df_MVBS = build_MVBS_ledger(df_Sv, slice_mins=20) - blocked = propagate_blocked_status( + blocked = propagate_status( df_Sv, df_MVBS, upstream_filename_column="raw_filename", @@ -295,6 +295,71 @@ def test_prediction_combines_two_aligned_mvbs_slices(): assert planned[0].filenames == ("first.zarr", "second.zarr") +def test_prediction_uses_completed_mvbs_and_skips_no_data_slice(): + mvbs = pd.DataFrame( + { + "MVBS_filename": ["first.zarr", "empty.zarr"], + "slice_start": pd.to_datetime(["2025-06-11T00:00:00Z", "2025-06-11T00:20:00Z"]), + "slice_end": pd.to_datetime(["2025-06-11T00:20:00Z", "2025-06-11T00:40:00Z"]), + "is_partial": [False, pd.NA], + "MVBS_status": ["completed", "no_data"], + } + ) + + prediction = build_prediction_ledger(mvbs, 40) + planned = plan_prediction_slices(mvbs, prediction) + + assert len(planned) == 1 + assert planned[0].filenames == ("first.zarr",) + assert planned[0].is_partial + + +def test_prediction_ledger_marks_all_no_data_window_terminal(): + starts = pd.to_datetime(["2025-06-11T00:00:00Z", "2025-06-11T00:20:00Z"]) + mvbs = pd.DataFrame( + { + "MVBS_filename": ["first.zarr", "second.zarr"], + "slice_start": starts, + "slice_end": starts + pd.Timedelta(minutes=20), + "is_partial": [pd.NA, pd.NA], + "MVBS_status": ["no_data", "no_data"], + } + ) + + prediction = build_prediction_ledger(mvbs, 40) + + assert prediction.loc[0, "prediction_status"] == "no_data" + assert plan_prediction_slices(mvbs, prediction) == [] + + +def test_existing_prediction_ledger_transitions_to_no_data(): + starts = pd.to_datetime(["2025-06-11T00:00:00Z", "2025-06-11T00:20:00Z"]) + mvbs = pd.DataFrame( + { + "MVBS_filename": ["first.zarr", "second.zarr"], + "slice_start": starts, + "slice_end": starts + pd.Timedelta(minutes=20), + "is_partial": [pd.NA, pd.NA], + "MVBS_status": ["pending", "pending"], + } + ) + prediction = build_prediction_ledger(mvbs, 40) + mvbs["MVBS_status"] = "no_data" + + updated = propagate_status( + mvbs, + prediction, + upstream_filename_column="MVBS_filename", + upstream_status_column="MVBS_status", + downstream_filenames_column="MVBS_filenames", + downstream_status_column="prediction_status", + upstream_label="MVBS slices", + ) + + assert updated.loc[0, "prediction_status"] == "no_data" + assert updated.loc[0, "error"] == "All required MVBS slices have no data" + + def test_prediction_planner_skips_completed_prediction_windows(): mvbs = pd.DataFrame( { @@ -315,6 +380,22 @@ def test_prediction_planner_skips_completed_prediction_windows(): assert planned == [] +def test_prediction_planner_skips_always_failed_prediction_windows(): + mvbs = pd.DataFrame( + { + "MVBS_filename": ["first.zarr", "second.zarr"], + "slice_start": pd.to_datetime(["2025-06-11T00:00:00Z", "2025-06-11T00:20:00Z"]), + "slice_end": pd.to_datetime(["2025-06-11T00:20:00Z", "2025-06-11T00:40:00Z"]), + "is_partial": [False, False], + "MVBS_status": ["completed", "completed"], + } + ) + prediction = build_prediction_ledger(mvbs, 40) + prediction.loc[0, ["prediction_status", "attempt_count"]] = ["always_failed", 3] + + assert plan_prediction_slices(mvbs, prediction) == [] + + def test_terminal_mvbs_failure_blocks_dependent_prediction_window(): starts = pd.date_range("2025-06-11T00:00:00Z", periods=4, freq="20min") mvbs = pd.DataFrame( @@ -330,7 +411,7 @@ def test_terminal_mvbs_failure_blocks_dependent_prediction_window(): ) prediction = build_prediction_ledger(mvbs, 40) - blocked = propagate_blocked_status( + blocked = propagate_status( mvbs, prediction, upstream_filename_column="MVBS_filename", @@ -394,6 +475,7 @@ def test_prediction_ledger_supports_overlapping_seven_minute_mvbs_slices(): "MVBS_filename": [f"slice-{index}.zarr" for index in range(6)], "slice_start": starts, "slice_end": starts + pd.Timedelta(minutes=7), + "MVBS_status": ["completed"] * 6, } )