diff --git a/src/echodataflow/deployment/core.py b/src/echodataflow/deployment/core.py index b435e97f..472d792a 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 2eddfd3b..6ff877f8 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 = ( @@ -352,6 +354,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): @@ -491,8 +512,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 @@ -503,7 +524,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( @@ -530,8 +551,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, " @@ -565,7 +586,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, @@ -574,6 +594,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"), @@ -616,6 +637,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/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 493126df..79d68aad 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 @@ -10,12 +11,11 @@ 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 prefect.events import emit_event -from echodataflow.flows.flows_helper import deployment_already_running -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,15 +34,16 @@ from echodataflow.operations.operations_postprocessing import ( build_MVBS_ledger, build_Sv_ledger, + failure_state, + propagate_status, plan_mvbs_slices, + plan_Sv_cleanup, 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, @@ -64,7 +65,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] = [], @@ -82,21 +83,6 @@ def flow_raw2Sv( add_location: bool = True, add_splitbeam_angle: bool = False, ): - # 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 # Assemble paths path_Sv_zarr = Path(path_main) / "Sv" @@ -417,17 +403,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, - 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, + 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, @@ -436,12 +423,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 @@ -449,7 +437,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( @@ -459,7 +447,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: @@ -470,7 +464,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, @@ -481,19 +474,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 @@ -513,6 +505,9 @@ def flow_raw2Sv_postprocessing( "first_ping_time", "last_ping_time", "error", + "Sv_cleanup_status", + "Sv_deleted_at", + "Sv_cleanup_error", ], ] = [ result.filename_raw, @@ -521,28 +516,43 @@ 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) 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 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") -@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, + no_data_gap_hours: float = 3.0, + 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", file_MVBS_csv: str = "MVBS_files.csv", + remove_completed_Sv_files: bool = True, ) -> None: """Create preplanned MVBS slices after all required raw conversions finish.""" + logger = get_run_logger() file_Sv_csv = Path(path_main) / file_Sv_csv if not file_Sv_csv.exists(): @@ -556,42 +566,62 @@ 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( 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 + updated_MVBS = propagate_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 - 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 = [] @@ -629,8 +659,54 @@ 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 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/flows/flows_helper.py b/src/echodataflow/flows/flows_helper.py index b8a3064c..b4c57ab4 100644 --- a/src/echodataflow/flows/flows_helper.py +++ b/src/echodataflow/flows/flows_helper.py @@ -1,8 +1,7 @@ from pathlib import Path import datetime -from prefect import flow, get_client, runtime, task -from prefect.client.schemas.filters import FlowRunFilter +from prefect import flow @flow(timeout_seconds=600, log_prints=True) @@ -72,23 +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 diff --git a/src/echodataflow/flows/flows_predict_hake.py b/src/echodataflow/flows/flows_predict_hake.py index 5ac39548..b8622fb1 100644 --- a/src/echodataflow/flows/flows_predict_hake.py +++ b/src/echodataflow/flows/flows_predict_hake.py @@ -20,7 +20,9 @@ ) from echodataflow.operations.operations_postprocessing import ( build_prediction_ledger, + failure_state, plan_prediction_slices, + propagate_status, read_or_create_ledger, ) from echodataflow.operations.operations_predict_hake import ( @@ -251,6 +253,8 @@ def flow_predict_hake_postprocessing( path_main: str, 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, @@ -260,6 +264,7 @@ def flow_predict_hake_postprocessing( file_prediction_csv: str = "prediction_files.csv", ) -> None: """Predict all newly ready windows, combining aligned MVBS slices.""" + logger = get_run_logger() file_MVBS_csv = Path(path_main) / file_MVBS_csv if not file_MVBS_csv.exists(): @@ -286,11 +291,28 @@ def flow_predict_hake_postprocessing( prediction_slice_mins, ), ) + + # Make terminal MVBS failures explicit in downstream planning + updated_prediction = propagate_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, 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 @@ -361,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 4acec330..9b243a98 100644 --- a/src/echodataflow/operations/operations_postprocessing.py +++ b/src/echodataflow/operations/operations_postprocessing.py @@ -131,14 +131,24 @@ 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"] = "" + ledger["Sv_cleanup_status"] = "pending" + ledger["Sv_deleted_at"] = pd.NaT + ledger["Sv_cleanup_error"] = "" 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) @@ -153,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( @@ -162,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", @@ -171,13 +189,79 @@ 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", - "error": "", + "MVBS_status": "no_data" if is_long_gap else "pending", + "attempt_count": 0, + "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) +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_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: + """Propagate terminal blocked and all-no-data upstream states downstream.""" + updated = downstream_ledger.copy() + if upstream_ledger.empty or downstream_ledger.empty: + return updated + + # 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) + ) + for idx, row in updated.iterrows(): + if row[downstream_status_column] not in {"pending", "failed"}: + continue + 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 + + def build_prediction_ledger( ledger_MVBS: pd.DataFrame, prediction_slice_mins: int = 40, @@ -205,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"), @@ -217,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) @@ -259,7 +345,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) @@ -285,6 +371,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, @@ -302,7 +414,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) @@ -310,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( @@ -324,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 735c454c..128d672e 100644 --- a/src/echodataflow/utils/manifests.py +++ b/src/echodataflow/utils/manifests.py @@ -27,9 +27,13 @@ "raw_filename", "Sv_filename", "raw2Sv_status", + "attempt_count", "first_ping_time", "last_ping_time", "error", + "Sv_cleanup_status", + "Sv_deleted_at", + "Sv_cleanup_error", ] MVBS_COLUMNS_POSTPROCESSING = [ "MVBS_filename", @@ -40,6 +44,7 @@ "last_ping_time", "is_partial", "MVBS_status", + "attempt_count", "error", ] PREDICTION_COLUMNS_POSTPROCESSING = [ @@ -54,6 +59,7 @@ "first_ping_time", "last_ping_time", "prediction_status", + "attempt_count", "error", ] 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 cdff2bd4..ff1b7380 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", @@ -643,3 +657,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/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..7f36750a 100644 --- a/tests/test_operations_postprocessing.py +++ b/tests/test_operations_postprocessing.py @@ -8,8 +8,11 @@ build_prediction_ledger, build_MVBS_ledger, build_Sv_ledger, + failure_state, generate_aligned_windows, + propagate_status, plan_mvbs_slices, + plan_Sv_cleanup, plan_prediction_slices, select_contained_records, select_overlapping_records, @@ -131,6 +134,85 @@ 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_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") + + 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_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(): @@ -163,6 +245,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) @@ -192,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( { @@ -212,6 +380,54 @@ 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( + { + "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_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( { @@ -259,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, } ) diff --git a/tests/test_postprocessing_flow_validation.py b/tests/test_postprocessing_flow_validation.py index bac5c645..90cb6ae9 100644 --- a/tests/test_postprocessing_flow_validation.py +++ b/tests/test_postprocessing_flow_validation.py @@ -1,8 +1,14 @@ +import pandas as pd + 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, ) from echodataflow.flows.flows_predict_hake import flow_predict_hake_postprocessing @@ -15,7 +21,6 @@ def info(self, message): messages.append(message) monkeypatch.setattr(flows_acoustics, "get_run_logger", lambda: Logger()) - flow_create_MVBS_postprocessing.fn(path_main=str(tmp_path)) assert messages == ["Sv ledger does not yet exist"] @@ -30,7 +35,6 @@ def info(self, message): messages.append(message) monkeypatch.setattr(flows_predict_hake, "get_run_logger", lambda: Logger()) - flow_predict_hake_postprocessing.fn( path_main=str(tmp_path), path_weight="unused.ckpt", @@ -38,3 +42,126 @@ 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(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"] + + +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"] + + +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"])