From 15cf43af69a1c5c90ab64be9deab76a94d371936 Mon Sep 17 00:00:00 2001 From: Carlos Osuna Date: Thu, 10 Sep 2026 14:38:20 +0200 Subject: [PATCH 01/21] enavle write inference to FDB --- config/varda-single-1.0-inference-only.yaml | 47 ++++ resources/fdb/realtime-varda.schema | 107 +++++++ ...-downscaler-global_trimedge_multi-fdb.yaml | 181 ++++++++++++ .../templates/templates_index_icon.yaml | 12 +- .../templates/varda-shortName=TOT_PREC.grib | Bin 0 -> 200 bytes .../varda-typeOfLevel=heightAboveGround.grib | Bin 0 -> 176 bytes .../varda-typeOfLevel=isobaricInhPa.grib | Bin 0 -> 176 bytes .../templates/varda-typeOfLevel=meanSea.grib | Bin 0 -> 176 bytes .../templates/varda-typeOfLevel=surface.grib | Bin 0 -> 176 bytes src/evalml/cli.py | 24 ++ src/evalml/config.py | 45 ++- workflow/Snakefile | 24 +- workflow/rules/common.smk | 26 +- workflow/rules/inference.smk | 265 +++++++++++++++++- workflow/rules/plot.smk | 6 +- workflow/rules/report.smk | 6 +- workflow/rules/verif_obs.smk | 12 +- workflow/rules/verification.smk | 30 +- workflow/scripts/inference_check_fdb.py | 196 +++++++++++++ .../scripts/inference_extract_requirements.py | 31 +- workflow/scripts/inference_prepare.py | 2 +- workflow/tools/config.schema.json | 111 ++++++-- 22 files changed, 1029 insertions(+), 96 deletions(-) create mode 100644 config/varda-single-1.0-inference-only.yaml create mode 100644 resources/fdb/realtime-varda.schema create mode 100644 resources/inference/configs/sgm-temporal-downscaler-global_trimedge_multi-fdb.yaml create mode 100644 resources/inference/templates/varda-shortName=TOT_PREC.grib create mode 100644 resources/inference/templates/varda-typeOfLevel=heightAboveGround.grib create mode 100644 resources/inference/templates/varda-typeOfLevel=isobaricInhPa.grib create mode 100644 resources/inference/templates/varda-typeOfLevel=meanSea.grib create mode 100644 resources/inference/templates/varda-typeOfLevel=surface.grib create mode 100644 workflow/scripts/inference_check_fdb.py diff --git a/config/varda-single-1.0-inference-only.yaml b/config/varda-single-1.0-inference-only.yaml new file mode 100644 index 00000000..650148c9 --- /dev/null +++ b/config/varda-single-1.0-inference-only.yaml @@ -0,0 +1,47 @@ +# yaml-language-server: $schema=../workflow/tools/config.schema.json +description: | + Run inference only for Varda-single-1.0 (no evaluation against observations). + +config_label: varda-single-1.0-inference-only + +dates: + start: 2025-03-01T00:00 + end: 2025-03-03T00:00 + frequency: 24h + +runs: + - temporal_downscaler: + checkpoint: https://service.meteoswiss.ch/mlstore#/models/sruc-m-2-interpolator/versions/3 + label: Varda-single-1.0 + steps: 0/120/1 + config: resources/inference/configs/sgm-temporal-downscaler-global_trimedge_multi-fdb.yaml + extra_requirements: + - anemoi-datasets==0.5.35 + - /scratch/mch/cosuna/evalml_fdb/anemoi-inference + forecaster: + checkpoint: https://service.meteoswiss.ch/mlstore#/models/sruc-m-1-forecaster/versions/4 + config: resources/inference/configs/sgm-multidataset-forecaster-global-ich1-oper.yaml + steps: 0/120/6 + +locations: + output_root: output/ + +profile: + fdb: + srun_uenv: "fdb/5.19:v2" + srun_view: "realtime" + # Base dir for the shared store; actual root per checkpoint is + # {fdb_root_global}/{env_id}. TODO: point at the real shared location. + fdb_root_global: "/scratch/mch/cosuna/evalml_fdb/global-fdb" + write_to_global_fdb: false + executor: slurm + global_resources: + gpus: 16 + default_resources: + slurm_partition: "postproc" + cpus_per_task: 1 + mem_mb_per_cpu: 1800 + runtime: "1h" + gpus: 0 + slurm_account: s83 + jobs: 50 diff --git a/resources/fdb/realtime-varda.schema b/resources/fdb/realtime-varda.schema new file mode 100644 index 00000000..c64fd7e4 --- /dev/null +++ b/resources/fdb/realtime-varda.schema @@ -0,0 +1,107 @@ + +# * Format of the rules is: + +# [a1, a2, a3 ...[b1, b2, b3... [c1, c2, c3...]]] + +# - The first level (a) defines which attributes are used to name the top level directory +# - The second level (b) defines which attributes are used to name the data files +# - The third level (c) defines which attributes are used as index keys + +# * Rules can be grouped + +# [a1, a2, a3 ... +# [b1, b2, b3... [c1, c2, c3...]] +# [B1, B2, B3... [C1, C2, C3...]] +# ] + +# * A list of values can be given for an attribute +# [ ..., stream=enfo/efov, ... ] +# This will be used when matching rules. + +# * Attributes can be typed +# Globally, at the begining of this file: + +# refdate: Date; + +# or in the context of a rule: +# [type=cl, ... [date:ClimateMonth, ...]] + +# Typing attributes is done when the user's requests or the GRIB values need to be modified before directories, files and indexes are created. For example, ClimateMonth will transform 2010-04-01 to 'may' internally. + +# * Attributes can be optional +# [ step, levelist?, param ] +# They will be replaced internally by an empty value. It is also posiible to provide a default subtitution value: e.g. [domain?g] will consider the domain to be 'g' if missing. + +# * Attributes can be removed: +# [grid-] +# This is useful to remove attributes present in the GRIB that should not be ignored + +# * Rules are matched: + +# - If the attributes are present in the GRIB/Request, or marked optional or ignored +# - If a list of possible value is provided, one of them must match, for example +# [ class, expver, stream=enfo/efov, date, time, domain ] +# will match either stream=enfo or stream=efov, all other attributes will be matched if they exist in the GRIB or user's request + +# * On archive: +# - Attributes are extracted from the GRIB (namespace 'mars'), possibly modified by the attribute type +# - Only the first rule is used, so order is important +# - All GRIB attributes must be used by the rules, otherwise an error is raised + +# * On retrieve: +# - Attributes are extracted from the user's request, possibly modified by the attribute type (e.g. for handling of U/V) +# - All the matching rules are considered +# - Only attributes listed in the rules are used to extract values from the user's request + + +# Default types + +param: Param; +step: Step; +date: Date; +hdate: Date; +refdate: Date; +latitude: Double; +longitude: Double; +level: Double; +levelist: Double; +grid: Grid; +expver: Expver; + +time: Time; +fcmonth: Integer; + +number: Integer; +frequency: Integer; +direction: Integer; +channel: Integer; + +instrument: Integer; +ident: Integer; + +diagnostic: Integer; +iteration: Integer; +system: Integer; +method: Integer; + +# ??????? + +# reference: Integer; +# fcperiod: Integer; + +# opttime: Integer; +# leadtime: Integer; + +# quantile: ?????? +# range: ?????? + +# band: Integer; + + +######################################################## +# These rules must be first, otherwise fields of These +# classes will be index with the default rule for oper #expver?, model?, type? + +[ date, time, stream, class, expver, model, type, domain- + [ levtype, number? + [ step, param, levelist?, timespan?none ]]] diff --git a/resources/inference/configs/sgm-temporal-downscaler-global_trimedge_multi-fdb.yaml b/resources/inference/configs/sgm-temporal-downscaler-global_trimedge_multi-fdb.yaml new file mode 100644 index 00000000..dc01a3ac --- /dev/null +++ b/resources/inference/configs/sgm-temporal-downscaler-global_trimedge_multi-fdb.yaml @@ -0,0 +1,181 @@ +runner: temporal_downscaler + +input: + cutout: + - lam_0: + grib: + path: forecaster/20* + pre_processors: + - forward_transform_filter: + # does not have an effect if the temporal downscaler does not use tp as prognostic variable + rescale: + scale: 0.001 # convert from kg m-2 to m + offset: 0 + param: tp + namer: &namer + rules: + # ICON short names (uppercase) + - - shortName: T + - t_{level} + - - shortName: U + - u_{level} + - - shortName: V + - v_{level} + - - shortName: W + - w_{level} + - - shortName: QV + - q_{level} + - - shortName: FI + - z_{level} + - - shortName: PMSL + - msl + - - shortName: FIS + - z + - - shortName: PS + - sp + - - shortName: T_2M + - 2t + - - shortName: TD_2M + - 2d + - - shortName: T_G + - skt + - - shortName: U_10M + - 10u + - - shortName: V_10M + - 10v + - - shortName: FR_LAND + - lsm + - - shortName: TOT_PREC + - tp + # IFS short names (lowercase) for the global grid output + - - shortName: t + levtype: pl + - t_{level} + - - shortName: u + levtype: pl + - u_{level} + - - shortName: v + levtype: pl + - v_{level} + - - shortName: w + levtype: pl + - w_{level} + - - shortName: q + levtype: pl + - q_{level} + - - shortName: z + levtype: pl + - z_{level} + - - shortName: z + levtype: sfc + - z + - - shortName: 10u + - 10u + - - shortName: 10v + - 10v + - - shortName: 2t + - 2t + - - shortName: 2d + - 2d + - - shortName: msl + - msl + - - shortName: sp + - sp + - - shortName: lsm + - lsm + - - shortName: skt + - skt + - - shortName: tp + - tp + - global: + grib: + path: forecaster/ifs* + namer: *namer + +constant_forcings: + test: + use_original_paths: true + +patch_metadata: resources/sgm-temporal-downscaler-ich1-oper-patch.yaml + +post_processors: + - accumulate_from_start_of_forecast: # accumulate tp from start of forecast + accumulations: + - tp + - forward_transform_filter: + rescale: + scale: 1000 # convert units from m to kg m-2 + offset: 0 + param: tp + +output: + tee: + - fdb: + encoding: + typeOfGeneratingProcess: 2 + generatingProcessIdentifier: 180 + # localNumberOfExperiment: 900 # not supported in current templates; marsClass=od fallback + templates: + samples: resources/templates_index_icon.yaml + post_processors: + - extract_mask: # removes global points + mask: "lam_0/cutout_mask" + as_slice: true + + - fdb: + encoding: + typeOfGeneratingProcess: 2 + generatingProcessIdentifier: 180 + # localNumberOfExperiment is not available in IFS/ECMWF GRIB templates; + # omitting it means marsClass=od (GPI=180 fallback) for the global grid + # Set native marsStream/marsType codetable keys so g2_mars_labeling returns + # stream=enfo and type=fc. Cannot use a concept override in local.98.def + # for marsStream because it conflicts with the codetable key that + # g2_mars_labeling expects as input. Setting them natively is portable + # across eccodes builds. + marsStream: enfo + marsType: fc + templates: + samples: resources/templates_index_ifs.yaml + post_processors: + - extract_mask: # removes lam points + mask: "lam_0/cutout_mask" + as_slice: true + inverse: true + - assign_mask: # fill local/global overlapping points with nan + mask: "global/cutout_mask" + modifiers: + - patches: + - variable: + U_10M: {"param": 165, "shortName": "10u"} + 10u: {"param": 165, "shortName": "10u"} + V_10M: {"param": 166, "shortName": "10v"} + 10v: {"param": 166, "shortName": "10v"} + TD_2M: {"param": 168, "shortName": "2d"} + 2d: {"param": 168, "shortName": "2d"} + T_2M: {"param": 167, "shortName": "2t"} + 2t: {"param": 167, "shortName": "2t"} + FR_LAND: {"param": 172, "shortName": "lsm"} + lsm: {"param": 172, "shortName": "lsm"} + PMSL: {"param": 151, "shortName": "msl"} + msl: {"param": 151, "shortName": "msl"} + PS: {"param": 134, "shortName": "sp"} + sp: {"param": 134, "shortName": "sp"} + SSO_SIGMA: {"param": 163, "shortName": "slor"} + slor: {"param": 163, "shortName": "slor"} + SSO_STDH: {"param": 160, "shortName": "sdor"} + sdor: {"param": 160, "shortName": "sdor"} + TOT_PREC: {"param": 228, "shortName": "tp"} + tp: {"param": 228, "shortName": "tp"} + z: {"param": 129, "shortName": "z"} + "^q_(\\d+)$": {"param": 133, "shortName": "q"} + "^t_(\\d+)$": {"param": 130, "shortName": "t"} + "^u_(\\d+)$": {"param": 131, "shortName": "u"} + "^v_(\\d+)$": {"param": 132, "shortName": "v"} + "^w_(\\d+)$": {"param": 135, "shortName": "w"} + "^z_(\\d+)$": {"param": 129, "shortName": "z"} + +# silenced due to bug in anemoi-inference for multi-step temporal downscalers, can be removed when fixed +verbosity: 0 +allow_nans: true +output_frequency: "1h" diff --git a/resources/inference/templates/templates_index_icon.yaml b/resources/inference/templates/templates_index_icon.yaml index b8861d22..67667c69 100644 --- a/resources/inference/templates/templates_index_icon.yaml +++ b/resources/inference/templates/templates_index_icon.yaml @@ -1,16 +1,16 @@ -# ICON-CH1 templates +# VARDA templates (PDT=0/8, typeOfGeneratingProcess=2 — deterministic forecast) - - {levtype: pl, param: [T, QV, U, V, OMEGA, FI]} - - resources/icon-ch1-typeOfLevel=isobaricInhPa.grib + - resources/varda-typeOfLevel=isobaricInhPa.grib - - {levtype: sfc, param: [T_2M, TD_2M, U_10M, V_10M]} - - resources/icon-ch1-typeOfLevel=heightAboveGround.grib + - resources/varda-typeOfLevel=heightAboveGround.grib - - {levtype: sfc, param: [FR_LAND, PS, FIS, T_G, SKT, SSO_SIGMA,SSO_STDH]} - - resources/icon-ch1-typeOfLevel=surface.grib + - resources/varda-typeOfLevel=surface.grib - - {levtype: sfc, param: [PMSL]} - - resources/icon-ch1-typeOfLevel=meanSea.grib + - resources/varda-typeOfLevel=meanSea.grib - - {levtype: sfc, param: [TOT_PREC]} - - resources/icon-ch1-shortName=TOT_PREC.grib + - resources/varda-shortName=TOT_PREC.grib diff --git a/resources/inference/templates/varda-shortName=TOT_PREC.grib b/resources/inference/templates/varda-shortName=TOT_PREC.grib new file mode 100644 index 0000000000000000000000000000000000000000..47deaf2d3db3d2b48b34ae81dbd98d8440f3c9e2 GIT binary patch literal 200 zcmZ<{@^t$DpMi-10!{!aQAURA4F8208QEVlG6MNvLXL^yF9S$`kwXEZl!2*^0VoaB z%)p?`%)lVn;sa!*vVk;+r`RqE_l(@;aw%l-**xt7KnW`rAi=?C!o<)EWPr?NVEm5& nFuTCwOkin{!5|>X3N~8{#BSgQk$ega4M2*G4WxvX-5dx2Y%C^~ literal 0 HcmV?d00001 diff --git a/resources/inference/templates/varda-typeOfLevel=heightAboveGround.grib b/resources/inference/templates/varda-typeOfLevel=heightAboveGround.grib new file mode 100644 index 0000000000000000000000000000000000000000..2aaa9491e360f6ad4dc1fba63395a7694c211b93 GIT binary patch literal 176 zcmZ<{@^t$DpMi-10yY3CQAURA4F8208QEVlG6MNvLXL^yF9S$`kwXrumZ^>bC=Jxi zz@W^`z#!P-17xMLfi#Gx*e(k9jNImODP-~4JnaKO2_+UF0R&78y$}WiLpqrM9}Iv@ dNmj5KS|HH|ZV<_*z|a7s*w{dJu(F#20RZ++ADjRH literal 0 HcmV?d00001 diff --git a/resources/inference/templates/varda-typeOfLevel=isobaricInhPa.grib b/resources/inference/templates/varda-typeOfLevel=isobaricInhPa.grib new file mode 100644 index 0000000000000000000000000000000000000000..334f77969c6486cd7231d5b2a949e1fa704d2a50 GIT binary patch literal 176 zcmZ<{@^t$DpMi-10yY3CQAURA4F8208QEVlG6MNvLXL^yF9S$`kwXxwmZ^>bC=Jxi zz@W^`z#!P-17xMLfi#Gx*e(k9jNImODP-~4JnaKO2_+UF0R&78y$}WiLrN1+?U(;x g02GsC1)HG-l5XGyk$ega4M2*G4P*x^yEza505%*U0{{R3 literal 0 HcmV?d00001 diff --git a/resources/inference/templates/varda-typeOfLevel=meanSea.grib b/resources/inference/templates/varda-typeOfLevel=meanSea.grib new file mode 100644 index 0000000000000000000000000000000000000000..f43f0748a72fc4d326fb9280c52130aef3802a73 GIT binary patch literal 176 zcmZ<{@^t$DpMi-10yY3CQAURA4F8208QEVlG6MNvLXL^yF9S$`kwX%ymZ^>bC=Jxi zz@W^`z#!P-17xMLfi#Gx*e(k9jNImODP-~4JnaKO2_+UF!NAPO#Lx?3fI#Yh1ORd+ cS;1y#futH(K_s67Lj#avV*}a2%5Dw>0OzD6i~s-t literal 0 HcmV?d00001 diff --git a/resources/inference/templates/varda-typeOfLevel=surface.grib b/resources/inference/templates/varda-typeOfLevel=surface.grib new file mode 100644 index 0000000000000000000000000000000000000000..71053eca5941929163830b32fac40e3a6cca77f8 GIT binary patch literal 176 zcmZ<{@^t$DpMi-10yY3CQAURA4F8208QEVlG6MNvLXL^yF9S$`kwX%ymZ^>bC=Jxi zz@W^`z#!P-17xMLfi#Gx*e(k9jNImODP-~4JnaKO2_+UF!N3f(sTafm0mlCb0OU%t ag3ZtZ%Yg|#1%?J7gN+Sj2P?Zd5C8zi93!Ft literal 0 HcmV?d00001 diff --git a/src/evalml/cli.py b/src/evalml/cli.py index f5c8a2cd..c1908f58 100644 --- a/src/evalml/cli.py +++ b/src/evalml/cli.py @@ -159,6 +159,30 @@ def cli(): pass +@cli.command(help="Run inference only (no verification) as defined by a config YAML file.") +@click.argument( + "configfile", type=click.Path(exists=True, dir_okay=False, path_type=Path) +) +@workflow_options +def inference( + configfile, cores, verbose, dry_run, unlock, report, dag, rulegraph, extra_smk_args +): + raise SystemExit( + execute_workflow( + configfile, + "inference_all", + cores, + verbose, + dry_run, + unlock, + report, + dag, + rulegraph, + extra_smk_args, + ) + ) + + @cli.command(help="Launch an experiment defined by a config YAML file.") @click.argument( "configfile", type=click.Path(exists=True, dir_okay=False, path_type=Path) diff --git a/src/evalml/config.py b/src/evalml/config.py index 3980046a..9730d7d3 100644 --- a/src/evalml/config.py +++ b/src/evalml/config.py @@ -66,6 +66,33 @@ class InferenceResources(BaseModel): None, description="Number of tasks per submission.", ) + srun_prefix: str | None = Field( + None, + description="Optional command to prepend to the srun invocation, e.g. 'uenv run --view=realtime fdb/5.19:v2 --'.", + ) + srun_uenv: str | None = Field( + None, + description="FDB uenv image to bundle into the venv squashfs, e.g. 'fdb/5.19:v2'. " + "When set, native FDB libraries are extracted from this uenv during venv creation.", + ) + srun_view: str | None = Field( + None, + description="FDB uenv view name to use when bundling FDB libraries, e.g. 'realtime'.", + ) + fdb_root_global: str | None = Field( + None, + description="Base directory for the shared/global FDB store. Each checkpoint gets " + "its own dedicated root at '{fdb_root_global}/{env_id}', keyed by the env_id hash " + "(derived from the checkpoint) so different checkpoints never mix data in the same " + "FDB instance. When set, inference is skipped for a given (run_id, init_time) if " + "this root already holds every requested lead time.", + ) + write_to_global_fdb: bool = Field( + False, + description="If true, inference writes directly into " + "'{fdb_root_global}/{env_id}' instead of the local per-checkpoint root. " + "Requires fdb_root_global to be set.", + ) class RunConfig(BaseModel): @@ -97,10 +124,6 @@ class RunConfig(BaseModel): description="List of extra dependencies to install for this model. " "These will be added to the requirements.txt file in the run directory.", ) - inference_resources: InferenceResources | None = Field( - None, - description="Resource requirements for inference jobs (optional; defaults handled externally).", - ) disable_local_eccodes_definitions: bool = Field( False, @@ -550,6 +573,10 @@ class Profile(BaseModel): """Workflow execution profile.""" executor: str = Field(..., description="Job executor, e.g. 'slurm'.") + fdb: InferenceResources | None = Field( + None, + description="FDB and inference resource settings (optional; overrides defaults for all inference jobs).", + ) global_resources: GlobalResources default_resources: DefaultResources jobs: int = Field(..., ge=1, description="Maximum number of parallel jobs.") @@ -671,14 +698,14 @@ class ConfigModel(BaseModel): ..., description="List of experiment participants, including forecaster/temporal downscaler ML runs and baselines.", ) - truth: TruthConfig | None + truth: TruthConfig | None = None lapse_rate_correction: bool = Field( default=True, description="Apply standard-atmosphere lapse-rate correction to T_2M.", ) - experiment: ExperimentConfig = Field( - ..., - description="Settings for the experiment workflow outputs.", + experiment: ExperimentConfig | None = Field( + None, + description="Settings for the experiment workflow outputs. Required for full experiment runs; optional for inference-only runs.", ) locations: Locations profile: Profile @@ -697,6 +724,8 @@ class ConfigModel(BaseModel): @model_validator(mode="after") def validate_scoremap_leadtimes(self) -> "ConfigModel": + if self.experiment is None: + return self sm = self.experiment.scoremaps if sm is None or not sm.enabled: return self diff --git a/workflow/Snakefile b/workflow/Snakefile index 67838e97..0734c239 100644 --- a/workflow/Snakefile +++ b/workflow/Snakefile @@ -3,6 +3,7 @@ # https://snakemake.readthedocs.io/en/stable/snakefiles/best_practices.html, import time +from types import SimpleNamespace from evalml.config import ConfigModel @@ -36,7 +37,8 @@ CONFIG_LABEL = config["config_label"] or CONFIG_FILE.stem EXPERIMENT_NAME = f"{WHEN}_{CONFIG_LABEL}_{CONFIG_HASH}" CANDIDATE_LABELS = ", ".join(v.get("label", k) for k, v in CANDIDATES.items()) BASELINE_LABELS = ", ".join(v.get("label", k) for k, v in BASELINES.items()) -GROUNDTRUTH_LABEL = config["truth"].get("label", config["truth"]["root"]) +_truth_cfg = config.get("truth") +GROUNDTRUTH_LABEL = _truth_cfg.get("label", _truth_cfg["root"]) if _truth_cfg else None DATA_DIR = OUT_ROOT / "data" LOGS_DIR = OUT_ROOT / "logs" @@ -134,7 +136,7 @@ onerror: # ----------------------------------------------------- -SCOREMAPS_CONFIGS = config.get("experiment", {}).get("scoremaps") or {} +SCOREMAPS_CONFIGS = (config.get("experiment") or {}).get("scoremaps") or {} rule experiment_all: @@ -160,7 +162,7 @@ rule experiment_all: scorecard_name=list(SCORECARD_CONFIGS), ) ] - if (config.get("experiment", {}).get("scorecards") or {}).get( + if (config.get("experiment") or {}).get("scorecards", {}) and (config.get("experiment") or {}).get("scorecards", {}).get( "enabled", True ) else [] @@ -235,14 +237,20 @@ rule sandbox_all: ), +def _inference_all_okfiles(wc): + return [ + _get_inference_okfile( + SimpleNamespace(run_id=run_id, init_time=t.strftime("%Y%m%d%H%M")) + ) + for run_id in CANDIDATES + for t in REFTIMES + ] + + rule inference_all: """Run inference for all reference times as defined in the configuration.""" input: - expand( - rules.inference_execute.output.okfile, - init_time=[t.strftime("%Y%m%d%H%M") for t in REFTIMES], - run_id=CANDIDATES, - ), + _inference_all_okfiles, output: run_all_ok=touch(OUT_ROOT / "logs/run_inference_all.ok"), shell: diff --git a/workflow/rules/common.smk b/workflow/rules/common.smk index acf76746..5342d795 100644 --- a/workflow/rules/common.smk +++ b/workflow/rules/common.smk @@ -19,7 +19,7 @@ ENV_HASH_FIELDS = { "disable_local_eccodes_definitions", } # Fields excluded from ALL hashing (display/resource metadata only). -RUN_HASH_EXCLUDE = {"label", "inference_resources", "_is_candidate", "model_type"} +RUN_HASH_EXCLUDE = {"label", "_is_candidate", "model_type", "fdb_root"} # Fields excluded from baseline hashing (display metadata only). BASELINE_HASH_EXCLUDE = {"label"} TRUTH_HASH_EXCLUDE = {"label"} @@ -88,6 +88,8 @@ def parse_regions(): """ from evalml.config import PREDEFINED_REGIONS + if config.get("experiment") is None: + return "[]" cfg = config["experiment"]["stratification"] root = cfg.get("root", "") result = [] @@ -388,8 +390,10 @@ def master_hash() -> str: return generate_json_hash(configs_to_hash) -def truth_hash(truth_config: dict) -> str: +def truth_hash(truth_config: dict | None) -> str: """Generate a short hash of the configs for the truth data.""" + if truth_config is None: + return generate_json_hash({}) cfg = {k: v for k, v in truth_config.items() if k not in TRUTH_HASH_EXCLUDE} return generate_json_hash(cfg) @@ -401,9 +405,11 @@ def verif_hash(full_config: dict) -> str: changing either (e.g. switching lapse_rate_correction on/off) produces new output paths and unconditionally triggers a rerun. """ - truth_cfg = { - k: v for k, v in full_config["truth"].items() if k not in TRUTH_HASH_EXCLUDE - } + truth_cfg = ( + {k: v for k, v in full_config["truth"].items() if k not in TRUTH_HASH_EXCLUDE} + if full_config.get("truth") is not None + else {} + ) experiment_verif_cfg = { "lapse_rate_correction": full_config.get("lapse_rate_correction", True), } @@ -413,6 +419,8 @@ def verif_hash(full_config: dict) -> str: def truth_file_dep(_): """Truth file dependency: a real path for zarr, but a live-query marker (no input file) for jretrieve.""" + if config.get("truth") is None: + return [] root = config["truth"]["root"] return [] if "jretrieve" in str(root) else [root] @@ -420,14 +428,14 @@ def truth_file_dep(_): # Fail fast: when the truth source is the live DWH (jretrievedwh), verify its # prerequisites at workflow-build time so a misconfigured environment is caught # at launch, before any (expensive) inference job runs. -if "jretrieve" in str(config["truth"]["root"]): +if config.get("truth") is not None and "jretrieve" in str(config["truth"]["root"]): from data_input.jretrieve import check_prerequisites, parse_selection _, _jretrieve_stage, _ = parse_selection(config["truth"]["root"]) check_prerequisites(_jretrieve_stage) -TRUTH_HASH = truth_hash(config["truth"]) +TRUTH_HASH = truth_hash(config.get("truth")) REGIONS = parse_regions() VERIF_HASH = verif_hash(config) _showcase = config.get("showcase", {}) @@ -436,7 +444,7 @@ SHOWCASE_CONFIG = { "params": _showcase.get("params", ["T_2M", "SP_10M"]), "fps": _showcase.get("animations", {}).get("frames_per_second", 2.0), } -EXPERIMENT_PARAMS = config.get("experiment", {}).get( +EXPERIMENT_PARAMS = (config.get("experiment") or {}).get( "params", ["T_2M", "TD_2M", "SP_10M", "PS", "PMSL", "TOT_PREC6"] ) REFTIMES = parse_reference_times() @@ -444,7 +452,7 @@ RUN_CONFIGS = collect_all_runs() ENV_CONFIGS = collect_all_envs() BASELINE_CONFIGS = collect_all_baselines() EXPERIMENT_PARTICIPANTS = collect_experiment_participants() -_scorecard = config.get("experiment", {}).get("scorecards") or {} +_scorecard = (config.get("experiment") or {}).get("scorecards") or {} SCORECARD_CONFIGS = ( _scorecard.get("sections", {}) if _scorecard.get("enabled", True) else {} ) diff --git a/workflow/rules/inference.smk b/workflow/rules/inference.smk index 1a83b57c..f035247f 100644 --- a/workflow/rules/inference.smk +++ b/workflow/rules/inference.smk @@ -94,6 +94,9 @@ rule inference_prepare_env: log: OUT_ROOT / "logs/inference_prepare_env/{env_id}.log", localrule: True + params: + fdb_uenv=lambda wc: _get_fdb_uenv_for_env(wc.env_id)[0], + fdb_view=lambda wc: _get_fdb_uenv_for_env(wc.env_id)[1], shell: """ ( @@ -109,6 +112,58 @@ rule inference_prepare_env: echo "[$(date)] Installing requirements from {input.requirements}..." uv pip install -r {input.requirements} + # Bundle FDB native libraries into the venv if this env needs FDB output. + # The FDB squashfs is mounted at /user-environment (its designed path) so that + # the absolute symlinks inside env/._// resolve correctly. + # The linux-zen3//lib64/ actual .so files are copied preserving their path + # structure, because libfdb5.so's RPATH is hardcoded to those absolute paths. + # The env/ symlink tree is copied as-is (preserved); after the venv squashfs + # is mounted at /user-environment at runtime, the full symlink chain resolves. + if [ -n "{params.fdb_uenv}" ]; then + echo "[$(date)] Bundling FDB native libraries from uenv '{params.fdb_uenv}' (view: {params.fdb_view})..." + FDB_SHA=$(uenv image inspect {params.fdb_uenv} | awk '/^sha:/ {{print $2}}') + FDB_REPO=$(uenv repo status | grep -oP '(?<=the repository at )\S+') + FDB_SQUASHFS="$FDB_REPO/images/$FDB_SHA/store.squashfs" + VENV_PATH=$VENV + FDB_VIEW={params.fdb_view} + + squashfs-mount "$FDB_SQUASHFS:/user-environment" -- bash -c ' + set -euo pipefail + VENV='"$VENV_PATH"' + VIEW='"$FDB_VIEW"' + + echo "[$(date)] Copying env/ view structure (symlinks preserved)..." + cp -rp /user-environment/env "$VENV/" + + echo "[$(date)] Copying linux-zen3 shared libraries (.so files, maintaining path structure)..." + find /user-environment/linux-zen3 -name "*.so*" -not -type l | while IFS= read -r f; do + rel="${{f#/user-environment/}}" + dest_dir="$VENV/${{rel%/*}}" + mkdir -p "$dest_dir" + cp -p "$f" "$dest_dir/" + done + echo "[$(date)] Copied $(find "$VENV/linux-zen3" -name "*.so*" 2>/dev/null | wc -l) shared libraries" + + echo "[$(date)] Copying FDB config and schema files..." + mkdir -p "$VENV/meta/recipe/meta/private/fdb_config" + cp "/user-environment/meta/recipe/meta/private/fdb_config/${{VIEW}}.yaml" \ + "$VENV/meta/recipe/meta/private/fdb_config/" + cp "/user-environment/meta/recipe/meta/private/fdb_config/${{VIEW}}.schema" \ + "$VENV/meta/recipe/meta/private/fdb_config/" 2>/dev/null || true + echo "[$(date)] FDB config files copied for view: $VIEW" + + echo "[$(date)] Cloning eccodes-cosmo-mars (varda-ext branch) for MARS definitions..." + mkdir -p "$VENV/data/share" + git clone --branch varda-ext --depth 1 \ + git@github.com:meteoswiss/eccodes-cosmo-mars.git \ + "$VENV/data/share/eccodes-cosmo-mars" + echo "[$(date)] Cloned eccodes-cosmo-mars definitions" + ' + + echo "[$(date)] Installing pyfdb (pure-Python FDB binding)..." + uv pip install pyfdb + fi + echo "[$(date)] Compiling Python bytecode..." python -m compileall -j 8 -o 0 -o 1 -o 2 $VENV/lib/python*/site-packages echo "[$(date)] Testing that eccodes is working..." @@ -155,14 +210,155 @@ rule inference_create_sandbox: def get_resource(wc, field: str, default): - """Fetch a resource field from the run config, or return the default.""" - rc = RUN_CONFIGS[wc.run_id] - if rc["inference_resources"] is None: - return default - if isinstance(rc["inference_resources"], dict): - return rc["inference_resources"].get(field, default) or default + """Fetch a resource field from profile.fdb, or return the default.""" + fdb = (config.get("profile") or {}).get("fdb") or {} + if isinstance(fdb, dict): + return fdb.get(field, default) or default + else: + return getattr(fdb, field, default) or default + + +def _get_fdb_uenv_for_env(env_id): + """Return (fdb_uenv, fdb_view) from profile.fdb.""" + fdb = (config.get("profile") or {}).get("fdb") or {} + if isinstance(fdb, dict): + uenv = fdb.get("srun_uenv", "") or "" + view = fdb.get("srun_view", "") or "realtime" + else: + uenv = getattr(fdb, "srun_uenv", "") or "" + view = getattr(fdb, "srun_view", "") or "realtime" + return uenv, view + + +def _get_fdb_roots(run_id): + """Return (global_root, local_root) for a run, both keyed by env_id. + + Each checkpoint (env_id) gets its own dedicated FDB instance, so different + checkpoints never mix data in the same root. global_root is '' if + profile.fdb.fdb_root_global is not configured. + """ + env_id = RUN_CONFIGS[run_id]["env_id"] + global_base = get_resource_by_env(env_id, "fdb_root_global", "") + global_root = f"{global_base}/{env_id}" if global_base else "" + local_root = str((OUT_ROOT / f"data/fdb/{env_id}").resolve()) + return global_root, local_root + + +def get_resource_by_env(env_id, field: str, default): + """Like get_resource, but usable outside a rule (no wildcards available).""" + fdb = (config.get("profile") or {}).get("fdb") or {} + if isinstance(fdb, dict): + return fdb.get(field, default) or default else: - return getattr(rc["inference_resources"], field) or default + return getattr(fdb, field, default) or default + + +def _count_fdb_blocks(run_id): + """Count distinct 'fdb:' targets in this run's own inference config output.tee. + + A single init_time can write several output blocks (e.g. an ICON-grid target + and an IFS-grid target) into the same env_id-keyed root; the check needs to + know how many to expect so a run that silently dropped one entirely doesn't + read as complete just because the other is. + """ + with open(RUN_CONFIGS[run_id]["config"]) as f: + run_cfg = yaml.safe_load(f) + tee = ((run_cfg.get("output") or {}).get("tee")) or [] + return sum(1 for entry in tee if isinstance(entry, dict) and "fdb" in entry) + + +checkpoint inference_check_fdb: + """Check whether (run_id, init_time) output already exists, complete, in FDB. + + Checks the global root first (if profile.fdb.fdb_root_global is set), then the + local per-checkpoint root, in that order. A 'hit' requires every requested lead + time to be present in one root -- see inference_check_fdb.py for why no further + MARS-key disambiguation (class/expver/model) is needed. Writes 'hit\\t' or + 'miss\\t' to the status file. Runs inside the squashfs venv so that pyfdb and the + native FDB C library are available. Skipped entirely (writes 'miss') when this + run's environment has no FDB uenv bundled (profile.fdb.srun_uenv unset). + """ + input: + image=lambda wc: OUT_ROOT + / f"data/runs/{RUN_CONFIGS[wc.run_id]['env_id']}/venv.squashfs", + script="workflow/scripts/inference_check_fdb.py", + output: + status=OUT_ROOT / "logs/inference_check_fdb/{run_id}-{init_time}.status", + log: + OUT_ROOT / "logs/inference_check_fdb/{run_id}-{init_time}.log", + localrule: True + params: + image_path=lambda wc, input: str(Path(input.image).resolve()), + fdb_configured=lambda wc: bool( + _get_fdb_uenv_for_env(RUN_CONFIGS[wc.run_id]["env_id"])[0] + ), + fdb_root_global=lambda wc: _get_fdb_roots(wc.run_id)[0], + fdb_root_local=lambda wc: _get_fdb_roots(wc.run_id)[1], + fdb_schema=str(Path("resources/fdb/realtime-varda.schema").resolve()), + fdb_view=lambda wc: get_resource(wc, "srun_view", "realtime") or "realtime", + date=lambda wc: datetime.strptime(wc.init_time, "%Y%m%d%H%M").strftime("%Y%m%d"), + time=lambda wc: datetime.strptime(wc.init_time, "%Y%m%d%H%M").strftime("%H%M"), + steps=lambda wc: RUN_CONFIGS[wc.run_id]["steps"], + expected_blocks=lambda wc: _count_fdb_blocks(wc.run_id), + shell: + """ + ( + set -euo pipefail + if [ "{params.fdb_configured}" != "True" ]; then + printf 'miss\t' > {output.status} + echo "FDB check skipped: FDB not configured for this environment" + else + GLOBAL_ARG="" + if [ -n "{params.fdb_root_global}" ]; then + GLOBAL_ARG="--fdb-root {params.fdb_root_global}" + fi + squashfs-mount {params.image_path}:/user-environment -- bash -c ' + source /user-environment/bin/activate + export LD_LIBRARY_PATH=/user-environment/env/{params.fdb_view}/lib64:/user-environment/env/{params.fdb_view}/lib${{LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}} + python {input.script} '"$GLOBAL_ARG"' \ + --fdb-root {params.fdb_root_local} \ + --fdb-schema {params.fdb_schema} \ + --date {params.date} \ + --time {params.time} \ + --steps {params.steps} \ + --expected-blocks {params.expected_blocks} \ + --output {output.status} + ' + fi + ) >{log} 2>&1 + """ + + +def _get_inference_status(wc): + """Return (okfile_path, fdb_root_or_None) for (run_id, init_time). + + 'hit' -> data already complete in FDB; okfile is the check's status file + (inference_execute is never requested), fdb_root points at the + root holding it. + 'miss' or FDB not configured -> falls through to inference_execute's okfile, + fdb_root is None (data will be local grib / freshly written FDB). + """ + env_id = RUN_CONFIGS[wc.run_id]["env_id"] + if _get_fdb_uenv_for_env(env_id)[0]: + status_file = checkpoints.inference_check_fdb.get( + run_id=wc.run_id, init_time=wc.init_time + ).output.status + status, _, root = Path(status_file).read_text().strip().partition("\t") + if status == "hit": + return str(status_file), root + return str(OUT_ROOT / f"logs/inference_execute/{wc.run_id}-{wc.init_time}.ok"), None + + +def _get_inference_okfile(wc): + """Dependency file marking (run_id, init_time) data as ready, whether freshly + executed or already complete in FDB.""" + return _get_inference_status(wc)[0] + + +def _get_inference_fdb_root(wc): + """FDB root already holding this (run_id, init_time)'s data (a check 'hit'), + or None if the data is local grib / about to be freshly executed.""" + return _get_inference_status(wc)[1] def get_leadtime(wc): @@ -217,7 +413,6 @@ rule inference_prepare_temporal_downscaler: output: config=Path(OUT_ROOT / "data/runs/{run_id}/{init_time}/config.yaml"), resources=directory(OUT_ROOT / "data/runs/{run_id}/{init_time}/resources"), - grib_out_dir=directory(OUT_ROOT / "data/runs/{run_id}/{init_time}/grib"), forecaster=directory(OUT_ROOT / "data/runs/{run_id}/{init_time}/forecaster"), okfile=touch( OUT_ROOT @@ -284,6 +479,18 @@ rule inference_execute: disable_local_definitions=lambda wc: RUN_CONFIGS[wc.run_id].get( "disable_local_eccodes_definitions", False ), + srun_prefix=lambda wc: get_resource(wc, "srun_prefix", "") or "", + fdb_view=lambda wc: get_resource(wc, "srun_view", "realtime") or "realtime", + fdb_configured=lambda wc: bool( + _get_fdb_uenv_for_env(RUN_CONFIGS[wc.run_id]["env_id"])[0] + ), + fdb_root=lambda wc: ( + _get_fdb_roots(wc.run_id)[0] + if get_resource(wc, "write_to_global_fdb", False) + and _get_fdb_roots(wc.run_id)[0] + else _get_fdb_roots(wc.run_id)[1] + ), + fdb_schema=str(Path("resources/fdb/realtime-varda.schema").resolve()), # fmt: off shell: """ @@ -292,6 +499,27 @@ rule inference_execute: cd {params.workdir} + # Write into this checkpoint's dedicated FDB root (env_id-keyed): the global + # one if write_to_global_fdb is set, otherwise the local per-checkpoint root. + # Exported so _run_inference (executed by squashfs-mount in a fresh bash) can see it. + export FDB5_CONFIG_OVERRIDE="" + if [ "{params.fdb_configured}" = "True" ]; then + mkdir -p "{params.fdb_root}" + # Use a patched schema that strips 'domain' (present in IFS GRIB MARS + # namespace via mars_labeling.def but absent from MeteoSwiss FDB schema). + cat > fdb5_config.yaml </lib64 symlink chain resolves to linux-zen3//lib64/*.so + # because both the env/ symlink tree and the actual .so files are in the venv. + if [ -d "$VENV/linux-zen3" ]; then + export LD_LIBRARY_PATH="$VENV/env/{params.fdb_view}/lib64:$VENV/env/{params.fdb_view}/lib${{LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}}" + # FDB5_CONFIG_OVERRIDE is set whenever FDB is configured for this env + # (see fdb_configured above); the bundled uenv config is only a fallback. + if [ -n "$FDB5_CONFIG_OVERRIDE" ]; then + export FDB5_CONFIG_FILE="$FDB5_CONFIG_OVERRIDE" + else + export FDB5_CONFIG_FILE="$VENV/meta/recipe/meta/private/fdb_config/{params.fdb_view}.yaml" + fi + # Prepend eccodes-cosmo-mars definitions so MARS namespace concepts + # (marsModel, marsStream, marsClass, etc.) resolve correctly for MeteoSwiss GRIBs. + # Must be set in ECCODES_DEFINITION_PATH (not just GRIB_DEFINITION_PATH) because + # the FDB C library uses eccodes internally to resolve MARS namespace keys. + export ECCODES_DEFINITION_PATH="$VENV/data/share/eccodes-cosmo-mars/definitions${{ECCODES_DEFINITION_PATH:+:$ECCODES_DEFINITION_PATH}}" + fi + CMD_ARGS=() # is GPU > 1, add parallel flag to CMD_ARGS and override automatic cluster detection @@ -307,7 +554,7 @@ rule inference_execute: CMD_ARGS+=(runner.parallel.cluster=slurm) fi - srun \ + {params.srun_prefix} srun \ --unbuffered \ --partition={resources.slurm_partition} \ --cpus-per-task={resources.cpus_per_task} \ diff --git a/workflow/rules/plot.smk b/workflow/rules/plot.smk index 5afdefd8..fe43605b 100644 --- a/workflow/rules/plot.smk +++ b/workflow/rules/plot.smk @@ -28,7 +28,7 @@ def _get_available_baselines(wc) -> list[dict[str, str]]: rule plot_meteogram: input: script="workflow/scripts/plot_meteogram.py", - inference_okfile=rules.inference_execute.output.okfile, + inference_okfile=_get_inference_okfile, truth_dep=truth_file_dep, output: expand( @@ -44,7 +44,7 @@ rule plot_meteogram: runtime="60m", params: ana_label=lambda wc: config["truth"]["label"], - truth_root=config["truth"]["root"], + truth_root=(config.get("truth") or {}).get("root", ""), fcst_grib=lambda wc: ( Path(OUT_ROOT) / f"data/runs/{wc.run_id}/{wc.init_time}/grib" ).resolve(), @@ -99,7 +99,7 @@ rule plot_meteogram: rule plot_forecast_frame: input: script="workflow/scripts/plot_forecast_frame.py", - inference_okfile=rules.inference_execute.output.okfile, + inference_okfile=_get_inference_okfile, output: expand( OUT_ROOT diff --git a/workflow/rules/report.smk b/workflow/rules/report.smk index 0ce78881..7106f6b3 100644 --- a/workflow/rules/report.smk +++ b/workflow/rules/report.smk @@ -9,7 +9,7 @@ include: "common.smk" def make_header_text(): dates = config["dates"] - truth = config["truth"]["label"] + truth = (config.get("truth") or {}).get("label", "") if isinstance(dates, list): return f"Explicit initializations from {len(dates)} runs have been used." text = f"Verification against {truth} with initializations from {dates.get('start')} to {dates.get('end')} by {dates.get('frequency')}" @@ -38,7 +38,7 @@ rule report_experiment_dashboard: params: sources=",".join(list(EXPERIMENT_PARTICIPANTS.keys())), header_text=make_header_text(), - stratification=" ".join(config["experiment"]["dashboard"]["stratification"]), + stratification=" ".join((config.get("experiment") or {}).get("dashboard", {}).get("stratification", [])), label_map=",".join( "{}:{}".format( sid, @@ -50,7 +50,7 @@ rule report_experiment_dashboard: ) for sid in EXPERIMENT_PARTICIPANTS ) - + ",truth-{}:{}".format(TRUTH_HASH, config["truth"]["label"]), + + ",truth-{}:{}".format(TRUTH_HASH, (config.get("truth") or {}).get("label", "")), shell: """ python {input.script} \ diff --git a/workflow/rules/verif_obs.smk b/workflow/rules/verif_obs.smk index c603870b..175a3964 100644 --- a/workflow/rules/verif_obs.smk +++ b/workflow/rules/verif_obs.smk @@ -1,6 +1,7 @@ import json from pathlib import Path from datetime import timedelta +from types import SimpleNamespace def _parse_steps(steps: str) -> list[int]: @@ -38,11 +39,12 @@ if config["mec"] is not None: rule prepare_mec_input: """Collect EKF SYNOP, monSYNOP, and reference verSYNOP observation files into the MEC input_obs directory.""" input: - inference_ok=lambda wc: expand( - rules.inference_execute.output.okfile, - run_id=wc.run_id, - init_time=[t.strftime("%Y%m%d%H%M") for t in REFTIMES], - ), + inference_ok=lambda wc: [ + _get_inference_okfile( + SimpleNamespace(run_id=wc.run_id, init_time=t.strftime("%Y%m%d%H%M")) + ) + for t in REFTIMES + ], output: obs=directory(OUT_ROOT / "data/runs/{run_id}/mec/{init_time}/input_obs"), ekf_file=OUT_ROOT diff --git a/workflow/rules/verification.smk b/workflow/rules/verification.smk index 425476ba..d0310f68 100644 --- a/workflow/rules/verification.smk +++ b/workflow/rules/verification.smk @@ -2,6 +2,7 @@ # VERIFICATION WORKFLOW # # ----------------------------------------------------- # from datetime import datetime +from types import SimpleNamespace import pandas as pd @@ -27,11 +28,11 @@ rule verification_metrics_baseline: params: baseline_steps=lambda wc: BASELINE_CONFIGS[wc.baseline_id]["steps"], member=lambda wc: BASELINE_CONFIGS[wc.baseline_id].get("member", "000"), - truth=config["truth"]["root"], + truth=(config.get("truth") or {}).get("root", ""), truth_source_id=f"truth-{TRUTH_HASH}", regions=REGIONS, experiment_params=",".join(EXPERIMENT_PARAMS), - threshold_dict=config["experiment"]["thresholds"], + threshold_dict=(config.get("experiment") or {}).get("thresholds", {}), lapse_rate_flag=( "--lapse_rate_correction" if config.get("lapse_rate_correction", True) @@ -68,7 +69,7 @@ rule verification_metrics: "src/verification/__init__.py", "src/data_input/__init__.py", script="workflow/scripts/verification_metrics.py", - inference_okfile=rules.inference_execute.output.okfile, + inference_okfile=_get_inference_okfile, truth_dep=truth_file_dep, output: OUT_ROOT / f"data/runs/{{run_id}}/{{init_time}}/verif_{VERIF_HASH}.nc", @@ -83,14 +84,14 @@ rule verification_metrics: # TODO: implement logic to use experiment name instead of run_id as wildcard params: fcst_steps=lambda wc: RUN_CONFIGS[wc.run_id]["steps"], - truth=config["truth"]["root"], + truth=(config.get("truth") or {}).get("root", ""), truth_source_id=f"truth-{TRUTH_HASH}", regions=REGIONS, grib_out_dir=lambda wc: ( Path(OUT_ROOT) / f"data/runs/{wc.run_id}/{wc.init_time}/grib" ).resolve(), experiment_params=",".join(EXPERIMENT_PARAMS), - threshold_dict=config["experiment"]["thresholds"], + threshold_dict=(config.get("experiment") or {}).get("thresholds", {}), lapse_rate_flag=( "--lapse_rate_correction" if config.get("lapse_rate_correction", True) @@ -185,7 +186,7 @@ rule verification_metrics_plot: ) for sid in EXPERIMENT_PARTICIPANTS ) - + ",truth-{}:{}".format(TRUTH_HASH, config["truth"]["label"]), + + ",truth-{}:{}".format(TRUTH_HASH, (config.get("truth") or {}).get("label", "")), shell: """ uv run {input.script} {input.verif} --output_dir {output} --labels "{params.label_map}" >{log} 2>&1 @@ -197,12 +198,13 @@ rule verification_scoremaps: "src/verification/__init__.py", "src/data_input/__init__.py", script="workflow/scripts/verification_scoremaps.py", - inference_okfiles=lambda wc: expand( - rules.inference_execute.output.okfile, - init_time=_restrict_reftimes_to_hours(REFTIMES), - allow_missing=True, - ), - truth=config["truth"]["root"], + inference_okfiles=lambda wc: [ + _get_inference_okfile( + SimpleNamespace(run_id=wc.run_id, init_time=it) + ) + for it in _restrict_reftimes_to_hours(REFTIMES) + ], + truth=(config.get("truth") or {}).get("root", ""), output: OUT_ROOT / f"data/runs/{{run_id}}/scoremaps/{{param}}_{{leadtime}}_{TRUTH_HASH}.nc", @@ -219,7 +221,7 @@ rule verification_scoremaps: params: fcst_label=lambda wc: RUN_CONFIGS[wc.run_id].get("label"), fcst_steps=lambda wc: RUN_CONFIGS[wc.run_id]["steps"], - truth_label=config["truth"]["label"], + truth_label=(config.get("truth") or {}).get("label", ""), reftimes=" ".join(t.strftime("%Y%m%d%H%M") for t in REFTIMES), run_root=lambda wc: (Path(OUT_ROOT) / f"data/runs/{wc.run_id}").resolve(), shell: @@ -242,7 +244,7 @@ rule verification_scoremaps_baseline: "src/data_input/__init__.py", script="workflow/scripts/verification_scoremaps.py", forecast=lambda wc: BASELINE_CONFIGS[wc.baseline_id]["root"], - truth=config["truth"]["root"], + truth=(config.get("truth") or {}).get("root", ""), output: OUT_ROOT / f"data/baselines/{{baseline_id}}/scoremaps/{{param}}_{{leadtime}}_{TRUTH_HASH}.nc", diff --git a/workflow/scripts/inference_check_fdb.py b/workflow/scripts/inference_check_fdb.py new file mode 100644 index 00000000..977ee639 --- /dev/null +++ b/workflow/scripts/inference_check_fdb.py @@ -0,0 +1,196 @@ +"""Check whether inference output for a given init_time already exists in FDB. + +Each FDB root passed on the command line is assumed dedicated to a single +checkpoint (the caller is responsible for pointing --fdb-root at an +env_id-keyed directory, e.g. '{fdb_root_global}/{env_id}' or +'{local_base}/{env_id}'), so no cross-checkpoint MARS-key disambiguation +(class/expver/model) is needed here. + +A single init_time's output can still be split into several distinct +*blocks* within that one root -- e.g. an ICON-grid target and an IFS-grid +target from the same run's `output.tee`, which go through different +post-processing and so don't necessarily share the same param set. A block +is identified by its combined key minus {step, param, levelist, timespan} +(i.e. by whatever class/stream/type/etc. distinguishes it). A "hit" requires: + + - at least `expected_blocks` distinct, complete blocks are present; + - each block holds every requested lead time (from --steps); + - each block's params are consistent across steps: every step other than + the first has exactly the same param set (the union across those steps), + and the first step's params are a subset of it. This tolerates fields + that are legitimately absent only at the initial step (e.g. accumulated + params with no meaningful value at lead time 0) without needing to name + them, while still catching a run that crashed partway through writing a + step's fields (which would otherwise still show up as "present"). + +Roots are checked in the order given; the first one holding a complete set +wins. Writes "hit\t" or "miss\t" to the output file. + +Intended to run inside a squashfs-mounted venv that provides pyfdb and the +native FDB C library (libfdb5.so). The squashfs is the same one built by +inference_create_venv / inference_make_squashfs_image. +""" +import argparse +import sys +from pathlib import Path + + +def _make_fdb_config(fdb_root: str, fdb_schema: str) -> dict: + return { + "type": "local", + "engine": "toc", + "schema": fdb_schema, + "spaces": [{"handler": "Default", "roots": [{"path": fdb_root}]}], + } + + +def _parse_steps(steps_spec: str) -> set[int]: + """Parse a 'start/end/step' lead-time spec (hours) into the expected step set.""" + start, end, step = map(int, steps_spec.split("/")) + return set(range(start, end + 1, step)) + + +def _group_by_block(entries) -> dict[tuple, dict[int, set]]: + """Group list() entries into {block_key: {step: {(levtype, param), ...}}}. + + block_key is the combined key with step/param/levelist/timespan/levtype + removed -- i.e. whatever identifies a distinct *output.tee target* + (class/stream/type/...) sharing this date+time. levtype is deliberately + NOT part of block_key: one target (e.g. an ICON-grid `fdb:` entry) writes + both 'sfc' and 'pl' levtypes, and those must be checked together as a + single block -- otherwise a run that crashed after writing only 'sfc' + would still count as its own separate, "complete" block, undermining the + --expected-blocks check. levtype is instead folded into each step's + param-identity tuple so sfc/pl params are never conflated with each other. + """ + blocks: dict[tuple, dict[int, set]] = {} + for el in entries: + key = el.combined_key() + if "step" not in key or "param" not in key: + continue + block_key = tuple( + sorted( + (k, v) + for k, v in key.items() + if k not in ("step", "param", "levelist", "timespan", "levtype") + ) + ) + item = (key.get("levtype"), key["param"]) + blocks.setdefault(block_key, {}).setdefault(int(key["step"]), set()).add(item) + return blocks + + +def _block_is_complete(step_params: dict[int, set], expected_steps: set[int]) -> bool: + found_steps = set(step_params) + if expected_steps - found_steps: + return False + + first_step = min(expected_steps) + other_steps = expected_steps - {first_step} + if not other_steps: + return True # single-step run: nothing to cross-check against + + reference = set() + for s in other_steps: + reference |= step_params[s] + if not reference: + return False # no params at all on the non-first steps + + for s in other_steps: + if step_params[s] != reference: + return False + return step_params[first_step] <= reference + + +def check_fdb( + fdb_root: str, + fdb_schema: str, + date: str, + time: str, + steps_spec: str, + expected_blocks: int, +) -> bool: + """Return True if `fdb_root` holds a complete output for this init time.""" + import pyfdb + + if not Path(fdb_root).is_dir(): + return False + + expected_steps = _parse_steps(steps_spec) + + # Pass the config as a dict directly. pyfdb.FDB(...) treats a `str` argument as + # inline YAML *content* to parse, not a file path (only a `Path` object is read + # as a file) -- passing str(config_path) here silently mis-parses and falls back + # to FDB5's built-in default config, which doesn't exist in this venv and fails + # with "Cannot open .../fdb5lib/etc/fdb/schema". A dict sidesteps the ambiguity + # entirely and needs no temp file. + with pyfdb.FDB(_make_fdb_config(fdb_root, fdb_schema)) as fdb: + # level=3 reaches the index level, where 'step'/'param' are part of the + # combined key (schema: [date,time,...,type[levtype,number? + # [step,param,levelist?,timespan?]]]). + entries = fdb.list({"date": date, "time": time}, level=3) + blocks = _group_by_block(entries) + + complete_blocks = [ + step_params + for step_params in blocks.values() + if _block_is_complete(step_params, expected_steps) + ] + if len(complete_blocks) < expected_blocks: + print( + f"FDB check: {fdb_root} has {len(complete_blocks)}/{expected_blocks} " + f"complete block(s) for date={date} time={time}" + ) + return False + return True + + +def main() -> int: + parser = argparse.ArgumentParser( + description="Check if FDB already holds a complete output for a given init_time." + ) + parser.add_argument( + "--fdb-root", + required=True, + action="append", + dest="fdb_roots", + help="Root path of an FDB store to check. Repeatable; checked in order, " + "first complete match wins.", + ) + parser.add_argument("--fdb-schema", required=True, help="Path to the FDB5 schema file.") + parser.add_argument("--date", required=True, help="Initialisation date (YYYYMMDD).") + parser.add_argument("--time", required=True, help="Initialisation time (HHMM).") + parser.add_argument( + "--steps", required=True, help="Expected lead times as 'start/end/step' (hours)." + ) + parser.add_argument( + "--expected-blocks", + required=True, + type=int, + help="Number of distinct fdb: output targets this run's config declares.", + ) + parser.add_argument("--output", required=True, help="File to write the result to.") + args = parser.parse_args() + + hit_root = "" + for root in args.fdb_roots: + if check_fdb( + root, args.fdb_schema, args.date, args.time, args.steps, args.expected_blocks + ): + hit_root = root + break + status = "hit" if hit_root else "miss" + + output = Path(args.output) + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text(f"{status}\t{hit_root}") + + print( + f"FDB check: date={args.date} time={args.time} steps={args.steps} " + f"-> {status}" + (f" ({hit_root})" if hit_root else "") + ) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/workflow/scripts/inference_extract_requirements.py b/workflow/scripts/inference_extract_requirements.py index dc3441ae..fcc85edf 100644 --- a/workflow/scripts/inference_extract_requirements.py +++ b/workflow/scripts/inference_extract_requirements.py @@ -180,6 +180,29 @@ def extract_git_requirements( return requirements +def _parse_local_package_name(path: str) -> str: + """Derive a package name from a local filesystem path. + + Reads the ``[project] name`` field from ``pyproject.toml`` if present; + falls back to the last path component. + """ + import os + import re + + pyproject = os.path.join(path, "pyproject.toml") + if os.path.exists(pyproject): + try: + with open(pyproject) as f: + content = f.read() + m = re.search(r'^\s*name\s*=\s*["\']([^"\']+)["\']', content, re.MULTILINE) + if m: + return m.group(1) + except Exception: + pass + + return os.path.basename(os.path.normpath(path)) + + def _parse_url_package_name(url: str) -> str: """Derive a package name from a VCS/HTTP URL. @@ -215,8 +238,11 @@ def parse_overrides(overrides: list[str]) -> dict[str, str | None]: - ``name==version`` → pinned PyPI package - a URL (``git+``, ``http://``, ``https://``) → VCS / direct install + - a local filesystem path (``/``, ``./``, ``../``) → local directory install - bare ``name`` → unpinned PyPI package (value is ``None``) """ + import os + result: dict[str, str | None] = {} for item in [*BASE_DEPENDENCIES, *(overrides or [])]: @@ -230,6 +256,9 @@ def parse_overrides(overrides: list[str]) -> dict[str, str | None]: elif any(item.startswith(prefix) for prefix in ("git+", "http://", "https://")): name = _parse_url_package_name(item) result[name] = item + elif item.startswith(("/", "./", "../")) or (os.sep in item and os.path.exists(item)): + name = _parse_local_package_name(item) + result[name] = item else: result[item] = None @@ -285,7 +314,7 @@ def format_requirements( pypi_requirements.pop(name, None) git_requirements.pop(name, None) - if isinstance(value, str) and value.startswith(("http://", "https://", "git+")): + if isinstance(value, str) and value.startswith(("http://", "https://", "git+", "/", "./", "../")): git_requirements[name] = value else: pypi_requirements[name] = value # type: ignore[assignment] # may be None diff --git a/workflow/scripts/inference_prepare.py b/workflow/scripts/inference_prepare.py index 59ba193d..766f2f52 100644 --- a/workflow/scripts/inference_prepare.py +++ b/workflow/scripts/inference_prepare.py @@ -46,7 +46,6 @@ def prepare_workdir(workdir: Path, resources_root: Path): Path to the root directory containing resource files. """ workdir.mkdir(parents=True, exist_ok=True) - (workdir / "grib").mkdir(parents=True, exist_ok=True) (workdir / "resources").mkdir(parents=True, exist_ok=True) shutil.copytree( resources_root / "templates", workdir / "resources", dirs_exist_ok=True @@ -116,6 +115,7 @@ def prepare_forecaster(smk): workdir = _get_workdir(smk) prepare_workdir(workdir, smk.params.resources_root) + (workdir / "grib").mkdir(parents=True, exist_ok=True) LOG.info("Prepared working directory at %s", workdir) res_list = "\n".join([str(fn) for fn in Path(workdir / "resources").rglob("*")]) LOG.info("Resources: \n%s", res_list) diff --git a/workflow/tools/config.schema.json b/workflow/tools/config.schema.json index ea2b84ba..7f85fa37 100644 --- a/workflow/tools/config.schema.json +++ b/workflow/tools/config.schema.json @@ -463,18 +463,6 @@ "title": "Extra Requirements", "type": "array" }, - "inference_resources": { - "anyOf": [ - { - "$ref": "#/$defs/InferenceResources" - }, - { - "type": "null" - } - ], - "default": null, - "description": "Resource requirements for inference jobs (optional; defaults handled externally)." - }, "disable_local_eccodes_definitions": { "default": false, "description": "If true, the ECCODES_DEFINITION_PATH environment variable will not be set to the COSMO local definitions.", @@ -609,6 +597,64 @@ "default": null, "description": "Number of tasks per submission.", "title": "Tasks" + }, + "srun_prefix": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Optional command to prepend to the srun invocation, e.g. 'uenv run --view=realtime fdb/5.19:v2 --'.", + "title": "Srun Prefix" + }, + "srun_uenv": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "FDB uenv image to bundle into the venv squashfs, e.g. 'fdb/5.19:v2'. When set, native FDB libraries are extracted from this uenv during venv creation.", + "title": "Srun Uenv" + }, + "srun_view": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "FDB uenv view name to use when bundling FDB libraries, e.g. 'realtime'.", + "title": "Srun View" + }, + "fdb_root_global": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Base directory for the shared/global FDB store. Each checkpoint gets its own dedicated root at '{fdb_root_global}/{env_id}', keyed by the env_id hash (derived from the checkpoint) so different checkpoints never mix data in the same FDB instance. When set, inference is skipped for a given (run_id, init_time) if this root already holds every requested lead time.", + "title": "Fdb Root Global" + }, + "write_to_global_fdb": { + "default": false, + "description": "If true, inference writes directly into '{fdb_root_global}/{env_id}' instead of the local per-checkpoint root. Requires fdb_root_global to be set.", + "title": "Write To Global Fdb", + "type": "boolean" } }, "title": "InferenceResources", @@ -692,6 +738,18 @@ "title": "Executor", "type": "string" }, + "fdb": { + "anyOf": [ + { + "$ref": "#/$defs/InferenceResources" + }, + { + "type": "null" + } + ], + "default": null, + "description": "FDB and inference resource settings (optional; overrides defaults for all inference jobs)." + }, "global_resources": { "$ref": "#/$defs/GlobalResources" }, @@ -941,18 +999,6 @@ "title": "Extra Requirements", "type": "array" }, - "inference_resources": { - "anyOf": [ - { - "$ref": "#/$defs/InferenceResources" - }, - { - "type": "null" - } - ], - "default": null, - "description": "Resource requirements for inference jobs (optional; defaults handled externally)." - }, "disable_local_eccodes_definitions": { "default": false, "description": "If true, the ECCODES_DEFINITION_PATH environment variable will not be set to the COSMO local definitions.", @@ -1086,7 +1132,8 @@ { "type": "null" } - ] + ], + "default": null }, "lapse_rate_correction": { "default": true, @@ -1095,8 +1142,16 @@ "type": "boolean" }, "experiment": { - "$ref": "#/$defs/ExperimentConfig", - "description": "Settings for the experiment workflow outputs." + "anyOf": [ + { + "$ref": "#/$defs/ExperimentConfig" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Settings for the experiment workflow outputs. Required for full experiment runs; optional for inference-only runs." }, "locations": { "$ref": "#/$defs/Locations" @@ -1137,8 +1192,6 @@ "description", "dates", "runs", - "truth", - "experiment", "locations", "profile" ], From 5b7d9ec0a514d80badc9dc162b1ea14e18ebcbf7 Mon Sep 17 00:00:00 2001 From: Carlos Osuna Date: Thu, 10 Sep 2026 23:22:27 +0200 Subject: [PATCH 02/21] mount two uenvs --- workflow/rules/inference.smk | 96 ++++++++++++++---------------------- 1 file changed, 37 insertions(+), 59 deletions(-) diff --git a/workflow/rules/inference.smk b/workflow/rules/inference.smk index f035247f..8c3d91b2 100644 --- a/workflow/rules/inference.smk +++ b/workflow/rules/inference.smk @@ -83,7 +83,10 @@ rule inference_extract_requirements: # Prepare the inference environment for a specific checkpoint. The venv is built # in /dev/shm (RAM) to avoid heavy IO on the parallel filesystem, then squashed -# to a .squashfs image tracked directly as the rule output. +# to a .squashfs image tracked directly as the rule output. FDB native libraries +# are NOT bundled into this image: the FDB uenv's own squashfs is mounted +# alongside the venv squashfs at run time instead (see inference_execute / +# inference_check_fdb), since squashfs-mount can mount several images at once. # See https://docs.cscs.ch/guides/storage/#python-virtual-environments-with-uenv. rule inference_prepare_env: input: @@ -96,7 +99,6 @@ rule inference_prepare_env: localrule: True params: fdb_uenv=lambda wc: _get_fdb_uenv_for_env(wc.env_id)[0], - fdb_view=lambda wc: _get_fdb_uenv_for_env(wc.env_id)[1], shell: """ ( @@ -112,56 +114,19 @@ rule inference_prepare_env: echo "[$(date)] Installing requirements from {input.requirements}..." uv pip install -r {input.requirements} - # Bundle FDB native libraries into the venv if this env needs FDB output. - # The FDB squashfs is mounted at /user-environment (its designed path) so that - # the absolute symlinks inside env/._// resolve correctly. - # The linux-zen3//lib64/ actual .so files are copied preserving their path - # structure, because libfdb5.so's RPATH is hardcoded to those absolute paths. - # The env/ symlink tree is copied as-is (preserved); after the venv squashfs - # is mounted at /user-environment at runtime, the full symlink chain resolves. + # Only the pure-Python / non-uenv pieces are prepared here. Native FDB + # libraries, its view symlinks and its config/schema files stay in the FDB + # uenv's own squashfs, mounted alongside this venv at run time. if [ -n "{params.fdb_uenv}" ]; then - echo "[$(date)] Bundling FDB native libraries from uenv '{params.fdb_uenv}' (view: {params.fdb_view})..." - FDB_SHA=$(uenv image inspect {params.fdb_uenv} | awk '/^sha:/ {{print $2}}') - FDB_REPO=$(uenv repo status | grep -oP '(?<=the repository at )\S+') - FDB_SQUASHFS="$FDB_REPO/images/$FDB_SHA/store.squashfs" - VENV_PATH=$VENV - FDB_VIEW={params.fdb_view} - - squashfs-mount "$FDB_SQUASHFS:/user-environment" -- bash -c ' - set -euo pipefail - VENV='"$VENV_PATH"' - VIEW='"$FDB_VIEW"' - - echo "[$(date)] Copying env/ view structure (symlinks preserved)..." - cp -rp /user-environment/env "$VENV/" - - echo "[$(date)] Copying linux-zen3 shared libraries (.so files, maintaining path structure)..." - find /user-environment/linux-zen3 -name "*.so*" -not -type l | while IFS= read -r f; do - rel="${{f#/user-environment/}}" - dest_dir="$VENV/${{rel%/*}}" - mkdir -p "$dest_dir" - cp -p "$f" "$dest_dir/" - done - echo "[$(date)] Copied $(find "$VENV/linux-zen3" -name "*.so*" 2>/dev/null | wc -l) shared libraries" - - echo "[$(date)] Copying FDB config and schema files..." - mkdir -p "$VENV/meta/recipe/meta/private/fdb_config" - cp "/user-environment/meta/recipe/meta/private/fdb_config/${{VIEW}}.yaml" \ - "$VENV/meta/recipe/meta/private/fdb_config/" - cp "/user-environment/meta/recipe/meta/private/fdb_config/${{VIEW}}.schema" \ - "$VENV/meta/recipe/meta/private/fdb_config/" 2>/dev/null || true - echo "[$(date)] FDB config files copied for view: $VIEW" - - echo "[$(date)] Cloning eccodes-cosmo-mars (varda-ext branch) for MARS definitions..." - mkdir -p "$VENV/data/share" - git clone --branch varda-ext --depth 1 \ - git@github.com:meteoswiss/eccodes-cosmo-mars.git \ - "$VENV/data/share/eccodes-cosmo-mars" - echo "[$(date)] Cloned eccodes-cosmo-mars definitions" - ' - echo "[$(date)] Installing pyfdb (pure-Python FDB binding)..." uv pip install pyfdb + + echo "[$(date)] Cloning eccodes-cosmo-mars (varda-ext branch) for MARS definitions..." + mkdir -p "$VENV/data/share" + git clone --branch varda-ext --depth 1 \ + git@github.com:meteoswiss/eccodes-cosmo-mars.git \ + "$VENV/data/share/eccodes-cosmo-mars" + echo "[$(date)] Cloned eccodes-cosmo-mars definitions" fi echo "[$(date)] Compiling Python bytecode..." @@ -292,6 +257,7 @@ checkpoint inference_check_fdb: fdb_configured=lambda wc: bool( _get_fdb_uenv_for_env(RUN_CONFIGS[wc.run_id]["env_id"])[0] ), + fdb_uenv=lambda wc: _get_fdb_uenv_for_env(RUN_CONFIGS[wc.run_id]["env_id"])[0], fdb_root_global=lambda wc: _get_fdb_roots(wc.run_id)[0], fdb_root_local=lambda wc: _get_fdb_roots(wc.run_id)[1], fdb_schema=str(Path("resources/fdb/realtime-varda.schema").resolve()), @@ -312,8 +278,11 @@ checkpoint inference_check_fdb: if [ -n "{params.fdb_root_global}" ]; then GLOBAL_ARG="--fdb-root {params.fdb_root_global}" fi - squashfs-mount {params.image_path}:/user-environment -- bash -c ' - source /user-environment/bin/activate + # Mount the venv (pyfdb, python) and the FDB uenv (native libs, config, + # at its required /user-environment path) side by side in one call. + FDB_SQUASHFS=$(uenv image inspect --format='{{sqfs}}' {params.fdb_uenv}) + squashfs-mount {params.image_path}:/venv-environment "$FDB_SQUASHFS:/user-environment" -- bash -c ' + source /venv-environment/bin/activate export LD_LIBRARY_PATH=/user-environment/env/{params.fdb_view}/lib64:/user-environment/env/{params.fdb_view}/lib${{LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}} python {input.script} '"$GLOBAL_ARG"' \ --fdb-root {params.fdb_root_local} \ @@ -481,6 +450,7 @@ rule inference_execute: ), srun_prefix=lambda wc: get_resource(wc, "srun_prefix", "") or "", fdb_view=lambda wc: get_resource(wc, "srun_view", "realtime") or "realtime", + fdb_uenv=lambda wc: _get_fdb_uenv_for_env(RUN_CONFIGS[wc.run_id]["env_id"])[0], fdb_configured=lambda wc: bool( _get_fdb_uenv_for_env(RUN_CONFIGS[wc.run_id]["env_id"])[0] ), @@ -522,23 +492,23 @@ FDBEOF _run_inference() {{ local VENV=$1 + local FDB_MOUNT=${{2:-}} source "$VENV/bin/activate" if [ "{params.disable_local_definitions}" = "False" ]; then export ECCODES_DEFINITION_PATH="$VENV/share/eccodes-cosmo-resources/definitions" fi - # Set up FDB env vars if native libraries were bundled into this venv. - # The env//lib64 symlink chain resolves to linux-zen3//lib64/*.so - # because both the env/ symlink tree and the actual .so files are in the venv. - if [ -d "$VENV/linux-zen3" ]; then - export LD_LIBRARY_PATH="$VENV/env/{params.fdb_view}/lib64:$VENV/env/{params.fdb_view}/lib${{LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}}" + # Set up FDB env vars when the FDB uenv is mounted alongside this venv + # (as a second image at $FDB_MOUNT -- see the squashfs-mount call below). + if [ -n "$FDB_MOUNT" ]; then + export LD_LIBRARY_PATH="$FDB_MOUNT/env/{params.fdb_view}/lib64:$FDB_MOUNT/env/{params.fdb_view}/lib${{LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}}" # FDB5_CONFIG_OVERRIDE is set whenever FDB is configured for this env - # (see fdb_configured above); the bundled uenv config is only a fallback. + # (see fdb_configured above); the uenv's own config is only a fallback. if [ -n "$FDB5_CONFIG_OVERRIDE" ]; then export FDB5_CONFIG_FILE="$FDB5_CONFIG_OVERRIDE" else - export FDB5_CONFIG_FILE="$VENV/meta/recipe/meta/private/fdb_config/{params.fdb_view}.yaml" + export FDB5_CONFIG_FILE="$FDB_MOUNT/meta/recipe/meta/private/fdb_config/{params.fdb_view}.yaml" fi # Prepend eccodes-cosmo-mars definitions so MARS namespace concepts # (marsModel, marsStream, marsClass, etc.) resolve correctly for MeteoSwiss GRIBs. @@ -566,7 +536,15 @@ FDBEOF }} export -f _run_inference - squashfs-mount {params.env_path}:/user-environment -- bash -c '_run_inference /user-environment' + # Mount the venv and (if configured) the FDB uenv side by side in one + # call: the FDB uenv's absolute internal symlinks require it to be at + # /user-environment, so the venv takes a different mountpoint instead. + if [ "{params.fdb_configured}" = "True" ]; then + FDB_SQUASHFS=$(uenv image inspect --format='{{sqfs}}' {params.fdb_uenv}) + squashfs-mount {params.env_path}:/venv-environment "$FDB_SQUASHFS:/user-environment" -- bash -c '_run_inference /venv-environment /user-environment' + else + squashfs-mount {params.env_path}:/venv-environment -- bash -c '_run_inference /venv-environment' + fi ) >{log} 2>&1 touch {output.okfile} """ From dfd84e24488a9a6191e4bfefc89f421f8c2bad9c Mon Sep 17 00:00:00 2001 From: Carlos Osuna Date: Tue, 15 Sep 2026 13:23:16 +0200 Subject: [PATCH 03/21] remove the fdb env, not needed since all deps are in python --- config/varda-single-1.0-inference-only.yaml | 2 + pyproject.toml | 1 + src/data_input/__init__.py | 232 +++++++++++++++++--- src/plotting/compat.py | 34 ++- uv.lock | 85 +++++++ workflow/rules/inference.smk | 103 ++++++--- workflow/rules/plot.smk | 76 ++++++- workflow/rules/verification.smk | 111 ++++++++-- workflow/scripts/plot_forecast_frame.py | 92 +++++--- workflow/scripts/verification_scoremaps.py | 36 ++- 10 files changed, 646 insertions(+), 126 deletions(-) diff --git a/config/varda-single-1.0-inference-only.yaml b/config/varda-single-1.0-inference-only.yaml index 650148c9..583faca3 100644 --- a/config/varda-single-1.0-inference-only.yaml +++ b/config/varda-single-1.0-inference-only.yaml @@ -22,6 +22,8 @@ runs: checkpoint: https://service.meteoswiss.ch/mlstore#/models/sruc-m-1-forecaster/versions/4 config: resources/inference/configs/sgm-multidataset-forecaster-global-ich1-oper.yaml steps: 0/120/6 + extra_requirements: + - anemoi-inference==0.11.1 locations: output_root: output/ diff --git a/pyproject.toml b/pyproject.toml index 06c6c69a..15088d1e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,6 +26,7 @@ dependencies = [ "scores>=2.0.0", "eccodes>=2.44,<2.48", "eccodes-cosmo-resources-python==2.44.0.1", + "pyfdb", # for reading VARDA-SINGLE forecast output archived directly to FDB "earthkit-data>=1.0,<1.1", "earthkit-utils>=1.0,<1.1", "earthkit-plots>=1.0,<1.1", diff --git a/src/data_input/__init__.py b/src/data_input/__init__.py index 23e52dd9..cef9064e 100644 --- a/src/data_input/__init__.py +++ b/src/data_input/__init__.py @@ -1,6 +1,8 @@ +import io import json import logging import re +import tempfile import urllib.request from datetime import datetime, timedelta from pathlib import Path @@ -39,9 +41,26 @@ } _ICON_TO_IFS = {v: k for k, v in _IFS_TO_ICON.items()} +# FDB schema shared by archival (workflow/rules/inference.smk) and reads (this module). +# Root marker prefix: load_forecast_data(root=...) treats an "fdb:" string the +# same way it already treats "jretrievedwh:..." for truth data -- a marker selecting a +# non-filesystem source, here an FDB root, rather than a real local directory. +_FDB_SCHEMA = Path(__file__).resolve().parents[2] / "resources/fdb/realtime-varda.schema" +_FDB_ROOT_MARKER = "fdb:" +# VARDA-SINGLE archives two grids under distinct `model` values: "varda-single" +# (ICON-origin, regional/LAM domain) and "varda-single-g" (IFS-origin, global +# domain) -- they have different point counts and can't be merged into one +# dataset via xarray. _collect_ml_grib_files' local-file equivalent only ever +# reads the regional grid (its "20*.grib" glob excludes the "ifs-"-prefixed +# global files), so load_forecast_data's FDB path filters to the same model for +# parity. plotting.compat.load_state_from_fdb uses both (see _FDB_GLOBAL_MODEL), +# mirroring load_state_from_grib's local ifs-*.grib companion-file merge, since +# that one works with flat point lists rather than an aligned xarray grid. +_FDB_MODEL = "varda-single" +_FDB_GLOBAL_MODEL = "varda-single-g" + XARRAY_ENGINE_PROFILE = { "ensure_dims": ["z", "number", "step", "forecast_reference_time"], - "add_valid_time_coord": True, "global_attrs": [{"institution": "MeteoSwiss"}, {"Conventions": "CF-1.8"}], } @@ -488,25 +507,28 @@ def load_from_grib_file(file: str | list[str], sel_kwargs): def variable_name_profile( level_type: Literal[ - "height_above_ground_level", - "mean_sea", + "heightAboveGround", + "meanSea", "surface", - "pressure", - "entire_atmosphere", + "isobaricInhPa", + "atmosphere", ], ) -> dict[str, Any]: - """Resolve variable name profile based on the level type.""" - if level_type in [ - "height_above_ground_level", - "mean_sea", - "surface", - "entire_atmosphere", - ]: + """Resolve variable name profile based on the raw GRIB `typeOfLevel` value. + + Note: earlier versions of this profile used earthkit-data's composite + "vertical.level_type"/"parameter.variable" metadata keys (semantic names + like "pressure", template "{parameter.variable}_{vertical.level}"). Those + keys don't resolve at all with the currently pinned earthkit-data (they + raise KeyError / silently match nothing) -- this now uses the raw, + always-available `typeOfLevel`/`shortName`/`levelist` GRIB keys instead. + """ + if level_type in ["heightAboveGround", "meanSea", "surface", "atmosphere"]: return {} - elif level_type == "pressure": + elif level_type == "isobaricInhPa": return { "variable_key": "p_l", - "remapping": {"p_l": "{parameter.variable}_{vertical.level}"}, + "remapping": {"p_l": "{shortName}_{levelist}"}, } else: raise ValueError(f"Unsupported level type: {level_type}") @@ -516,13 +538,19 @@ def fieldlist_to_xarray(fieldlist) -> xr.Dataset: ds = xr.Dataset() if len(fieldlist) == 0: return ds - for level_type_group in fieldlist.group_by("vertical.level_type"): - # earthkit-data should return the group key...TODO: open issue? - level_type = level_type_group.get("vertical.level_type")[0] + for level_type_group in fieldlist.group_by("typeOfLevel"): + level_type = level_type_group[0].metadata("typeOfLevel") profile = XARRAY_ENGINE_PROFILE | variable_name_profile(level_type) _ds = level_type_group.to_xarray(**profile, allow_holes=True) + # add_valid_time_coord=True (earthkit's own convenience for this) isn't + # supported together with allow_holes=True on the currently pinned + # earthkit-data (raises NotImplementedError) -- compute it by hand instead. + if "forecast_reference_time" in _ds.coords and "step" in _ds.coords: + _ds = _ds.assign_coords( + valid_time=_ds["forecast_reference_time"] + _ds["step"] + ) ds = ds.merge( - _ds, compat="no_conflicts", combine_attrs="no_conflicts", join="outer" + _ds, compat="no_conflicts", combine_attrs="drop_conflicts", join="outer" ) return ds @@ -621,15 +649,146 @@ def _load_forecast_data_from_grib(files: list[Path], params: list[str]) -> xr.Da {p for p in params} | {_ICON_TO_IFS[p] for p in params if p in _ICON_TO_IFS} ) ds = load_from_grib_file(files, {"parameter.variable": params_extended}) + return _rename_ifs_to_icon(ds) - # Rename any IFS shortNames back to ICON names + +def _rename_ifs_to_icon(ds: xr.Dataset) -> xr.Dataset: + """Rename any IFS shortNames present in `ds` back to their ICON equivalents.""" ifs_rename = { ifs: icon for ifs, icon in _IFS_TO_ICON.items() if ifs in ds.data_vars } - if ifs_rename: - ds = ds.rename(ifs_rename) + return ds.rename(ifs_rename) if ifs_rename else ds - return ds + +def fdb_has_data(fdb_root: str, reftime: datetime, model: str = _FDB_MODEL) -> bool: + """Lightweight existence check: does `fdb_root` have any entry for this + date/time under `model`? + + Doesn't check completeness (see workflow/scripts/inference_check_fdb.py's + block/step/param accounting for that) -- for callers where completeness was + already established upstream (e.g. by a Snakemake dependency on this + reftime's inference_check_fdb/inference_execute okfile before this runs), + a bare presence check is enough to decide whether to route a reftime to + the FDB path at all. + """ + import pyfdb + + config = { + "type": "local", + "engine": "toc", + "schema": str(_FDB_SCHEMA.resolve()), + "spaces": [{"handler": "Default", "roots": [{"path": fdb_root}]}], + } + with pyfdb.FDB(config) as fdb: + entries = list( + fdb.list( + { + "date": reftime.strftime("%Y%m%d"), + "time": reftime.strftime("%H%M"), + "model": model, + }, + level=1, + ) + ) + return bool(entries) + + +def _fetch_fdb_grib_bytes( + fdb_root: str, reftime: datetime, steps: list[int], model: str | None = None +) -> bytes: + """Retrieve raw GRIB bytes from an FDB store for one date/time and lead times. + + `fdb.retrieve()` requires a single value per request for whatever keys + identify a distinct *block* (class/stream/type/expver/model/levtype) -- + unlike `fdb.list()`, which can wildcard across them -- so this lists first + to discover the blocks actually archived for this date/time/steps (e.g. one + per `output.tee` `fdb:` target), then issues one `retrieve()` per block and + concatenates their bytes (concatenated GRIB messages are themselves a valid + GRIB stream, readable as one file). + + `model` restricts discovery to that exact value (VARDA-SINGLE archives + multiple grids under distinct model values, see `_FDB_MODEL`/ + `_FDB_GLOBAL_MODEL`); pass `None` to include every archived model instead. + Returns `b""` if nothing matches -- callers decide whether that's an error. + """ + import pyfdb + + date = reftime.strftime("%Y%m%d") + time_ = reftime.strftime("%H%M") + config = { + "type": "local", + "engine": "toc", + "schema": str(_FDB_SCHEMA.resolve()), + "spaces": [{"handler": "Default", "roots": [{"path": fdb_root}]}], + } + # Keys identifying a distinct archived block; retrieve() needs a single value + # for each of these (list() can wildcard across them, retrieve() can't). + block_keys = ("class", "stream", "type", "expver", "model", "levtype") + list_request = {"date": date, "time": time_, "step": [str(s) for s in steps]} + if model is not None: + list_request["model"] = model + + with pyfdb.FDB(config) as fdb: + entries = list(fdb.list(list_request, level=3)) + if not entries: + return b"" + + blocks: dict[tuple, dict[str, set]] = {} + for entry in entries: + key = entry.combined_key() + block_key = tuple(key[bkey] for bkey in block_keys) + dims = blocks.setdefault(block_key, {}) + for k, v in key.items(): + if v not in (None, "") and k not in block_keys: + dims.setdefault(k, set()).add(v) + + buf = io.BytesIO() + for block_key, dims in blocks.items(): + request = dict(zip(block_keys, block_key)) + request.update({k: sorted(v) for k, v in dims.items()}) + handle = fdb.retrieve(request) + handle.open() + buf.write(handle.read()) + handle.close() + return buf.getvalue() + + +def _load_forecast_data_from_fdb( + fdb_root: str, reftime: datetime, steps: list[int], params: list[str] +) -> xr.Dataset: + """Load forecast data directly from an FDB store (no local GRIB files). + + External callers should use :func:`load_forecast_data` (root prefixed with + ``"fdb:"``), which handles param expansion, step expansion, IC synthesis, + and disaggregation, same as the GRIB-file path. Restricted to `_FDB_MODEL` + (the regional grid) for parity with `_collect_ml_grib_files`'s local-file + equivalent -- see the comment on `_FDB_MODEL`. + + The retrieved bytes are written to a temp GRIB file and reopened via the + existing local-file path (:func:`load_from_grib_file`) rather than decoded + in memory: earthkit-data's own `"fdb"` source has a bug (it never opens the + `pyfdb` `DataHandle` before reading -- raises `IOProblemError`), and + fieldlists built via `.sel()` / `FieldList.from_fields()` don't behave like + file-backed ones for metadata batch access. Writing to disk and reopening + as a genuine file sidesteps both and reuses the exact same, already-tested + GRIB-loading path used for local files. + """ + data = _fetch_fdb_grib_bytes(fdb_root, reftime, steps, model=_FDB_MODEL) + if not data: + raise FileNotFoundError( + f"No FDB data found for date={reftime:%Y%m%d} time={reftime:%H%M} " + f"model={_FDB_MODEL} in {fdb_root}" + ) + + with tempfile.NamedTemporaryFile(suffix=".grib") as tmp: + tmp.write(data) + tmp.flush() + params_extended = list( + {p for p in params} | {_ICON_TO_IFS[p] for p in params if p in _ICON_TO_IFS} + ) + ds = load_from_grib_file(tmp.name, {"parameter.variable": params_extended}) + + return _rename_ifs_to_icon(ds) def _jretrieve_df_to_xarray(df, short_names, catalog) -> xr.Dataset: @@ -1339,14 +1498,35 @@ def load_forecast_data( - Plain ``TOT_PREC`` is returned as cumulative-from-start without disaggregation Routing (in order): - 1. ``*.grib`` files present in *root* → :func:`_load_forecast_data_from_grib` + 1. ``root`` starts with ``"fdb:"`` → :func:`_load_forecast_data_from_fdb` + (inference output archived directly to FDB, no local GRIB at all) + 2. ``*.grib`` files present in *root* → :func:`_load_forecast_data_from_grib` (ML inference output) - 2. ``INCA`` in path parts → :func:`_load_INCA_baseline_from_netcdf` - 3. Otherwise → ICON operational archive (via :func:`_load_icon_baseline_from_grib`) + 3. ``INCA`` in path parts → :func:`_load_INCA_baseline_from_netcdf` + 4. Otherwise → ICON operational archive (via :func:`_load_icon_baseline_from_grib`) """ - root = Path(root) load_params = get_base_params(params) load_steps = get_steps(steps, params) + if str(root).startswith(_FDB_ROOT_MARKER): + LOG.info("Loading forecasts from FDB...") + fdb_root = str(root)[len(_FDB_ROOT_MARKER) :] + ds = _load_forecast_data_from_fdb(fdb_root, reftime, load_steps, load_params) + # Try to derive elevation from surface geopotential (FIS/z at step 0) + # before falling back to ICON grid constants lookup. + try: + fis_ds = _load_forecast_data_from_fdb(fdb_root, reftime, [0], ["FIS"]) + except FileNotFoundError: + fis_ds = None + if fis_ds is not None and "FIS" in fis_ds: + elevation = ekdv.geopotential_height_from_geopotential( + fis_ds["FIS"].isel(step=0, drop=True).squeeze(drop=True) + ) + ds = ds.assign_coords(elevation=elevation) + if "elevation" not in ds.coords: + ds = _try_assign_elevation(ds) + return _disaggregated_and_derived_params(ds, steps, params) + + root = Path(root) if any(root.glob("*.grib")): LOG.info("Loading forecasts from GRIB files...") ds = _load_forecast_data_from_grib( diff --git a/src/plotting/compat.py b/src/plotting/compat.py index 3f63313a..07ff390a 100644 --- a/src/plotting/compat.py +++ b/src/plotting/compat.py @@ -1,10 +1,12 @@ from datetime import datetime from pathlib import Path +import tempfile + import geopandas as gpd import numpy as np from shapely.geometry import MultiPoint -from data_input import load_from_grib_file +from data_input import load_from_grib_file, _fetch_fdb_grib_bytes, _FDB_MODEL, _FDB_GLOBAL_MODEL PARAMS_MAP = { @@ -108,6 +110,36 @@ def load_state_from_grib( return state +def load_state_from_fdb( + fdb_root: str, reftime: datetime, step: int, paramlist: list[str] | None = None +) -> dict[str, np.ndarray | dict[str, np.ndarray] | gpd.GeoSeries]: + """FDB equivalent of :func:`load_state_from_grib` for one lead time. + + Fetches the regional (`_FDB_MODEL`) and global (`_FDB_GLOBAL_MODEL`) grids + from FDB and writes them to temp files using the same naming convention + `load_state_from_grib` already expects locally (a primary file plus an + `ifs-.grib` companion in the same directory), then delegates to it + unchanged -- reusing its state-dict construction and global-grid + concatenation (it works with flat point lists, not an aligned xarray grid, + so the regional/global point-count mismatch that :func:`data_input. + load_forecast_data` sidesteps by dropping the global grid entirely isn't a + problem here). + """ + regional = _fetch_fdb_grib_bytes(fdb_root, reftime, [step], model=_FDB_MODEL) + if not regional: + raise FileNotFoundError( + f"No FDB data found for step={step} model={_FDB_MODEL} in {fdb_root}" + ) + global_ = _fetch_fdb_grib_bytes(fdb_root, reftime, [step], model=_FDB_GLOBAL_MODEL) + + with tempfile.TemporaryDirectory() as tmp_dir: + primary = Path(tmp_dir) / "frame.grib" + primary.write_bytes(regional) + if global_: + (Path(tmp_dir) / "ifs-frame.grib").write_bytes(global_) + return load_state_from_grib(primary, paramlist=paramlist) + + def load_state_from_raw( file: Path, paramlist: list[str] | None = None ) -> dict[str, np.ndarray | dict[str, np.ndarray]]: diff --git a/uv.lock b/uv.lock index eea30160..a50205e2 100644 --- a/uv.lock +++ b/uv.lock @@ -1273,6 +1273,7 @@ dependencies = [ { name = "mlflow" }, { name = "netcdf4" }, { name = "pydantic" }, + { name = "pyfdb" }, { name = "pyproj" }, { name = "pyzmq" }, { name = "scores" }, @@ -1305,6 +1306,7 @@ requires-dist = [ { name = "mlflow", specifier = ">=3.1.1" }, { name = "netcdf4", specifier = ">=1.7.2" }, { name = "pydantic", specifier = ">=2.11.7" }, + { name = "pyfdb" }, { name = "pyproj", specifier = ">=3.7.2" }, { name = "pyzmq", specifier = ">=27.1.0" }, { name = "scores", specifier = ">=2.0.0" }, @@ -1375,6 +1377,34 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/cb/a8/20d0723294217e47de6d9e2e40fd4a9d2f7c4b6ef974babd482a59743694/fastjsonschema-2.21.2-py3-none-any.whl", hash = "sha256:1c797122d0a86c5cace2e54bf4e819c36223b552017172f32c5c024a6b77e463", size = 24024, upload-time = "2025-08-14T18:49:34.776Z" }, ] +[[package]] +name = "fdb5lib" +version = "5.21.4.21" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "eccodeslib" }, + { name = "eckitlib" }, + { name = "metkitlib" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/21/42/ad862302f4252ee81045a1d3b7dfd18c91afe71a565c2009ed3ead4e2828/fdb5lib-5.21.4.21-cp311-cp311-macosx_13_0_arm64.whl", hash = "sha256:47457a9e524adc1e6386af29fcfdd8134e9ce302918eb88dcd40687c9bbef52b", size = 1383780, upload-time = "2026-06-25T14:50:12.38Z" }, + { url = "https://files.pythonhosted.org/packages/1b/0d/225080639df23bbc168afbe8a315713127946581650c38104f392dd4369b/fdb5lib-5.21.4.21-cp311-cp311-macosx_13_0_x86_64.whl", hash = "sha256:c8382a26ddb024750e0ef2079900c37bbb85e0a714577930d0f6eb964a6c9792", size = 1413311, upload-time = "2026-06-25T14:39:14.826Z" }, + { url = "https://files.pythonhosted.org/packages/0c/d1/08c3a597684b40901c39cd932c7e3a11590f85f2d9c8545472ac70754223/fdb5lib-5.21.4.21-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:62a31a0f9d65ff97d87c2c9e2dd966a68f68fd34f048ee4c8c906f69e942ff89", size = 1687866, upload-time = "2026-06-25T14:35:12.348Z" }, + { url = "https://files.pythonhosted.org/packages/f6/67/3f6fb9e06f5434e3c001881eae57ea1e1fe33602cf1ed0375fdce343de44/fdb5lib-5.21.4.21-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:3b1f97c93c02d1eac2a0f7c80207942fc7ee1687117c90bf9c31b75957f38f4f", size = 1667418, upload-time = "2026-06-25T14:35:24.423Z" }, + { url = "https://files.pythonhosted.org/packages/30/b6/7cf7763e593c920b54e7be47b0a676a819b5b2aa7416e083444585eae04f/fdb5lib-5.21.4.21-cp312-cp312-macosx_13_0_arm64.whl", hash = "sha256:4328703222058fe67dc590c7957aa0b09d92887455212c0bfd851fed2edad4bb", size = 1383779, upload-time = "2026-06-25T15:07:59.25Z" }, + { url = "https://files.pythonhosted.org/packages/ce/5f/a36d1c694859bce2511b548fea4fb42175689fb44e96d999ba9f6e1c5af7/fdb5lib-5.21.4.21-cp312-cp312-macosx_13_0_x86_64.whl", hash = "sha256:fe343e453db17ae80adad14d13f169feb99a7a81baa0446cb9dc6addba4ec5eb", size = 1413310, upload-time = "2026-06-25T15:20:02.611Z" }, + { url = "https://files.pythonhosted.org/packages/b9/2e/6168244bf0da56887fb7a212276326f18b8a7553378e0d20e7f4c29f6fc0/fdb5lib-5.21.4.21-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:6d072697b93ebfda93791993a68188baee8137ea7a5fbd540333a85d80a0b9ec", size = 1687865, upload-time = "2026-06-25T14:36:11.968Z" }, + { url = "https://files.pythonhosted.org/packages/13/9a/c716cc83f3557d678de74acd7c9dbad3eeacbde227fda9d7db8bea5fbe0a/fdb5lib-5.21.4.21-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:1bd16c7515840aff8fcdf4fdb260a7c23e6d1e8c9ea54e0c7cea3a5974eba16a", size = 1667418, upload-time = "2026-06-25T14:35:51.007Z" }, + { url = "https://files.pythonhosted.org/packages/66/48/149ef9a854187423dbdf57133e2ba3655d6a784acf1bf51b9d56903c1dd8/fdb5lib-5.21.4.21-cp313-cp313-macosx_13_0_arm64.whl", hash = "sha256:afd95eab549646717c5258757661c7b30854b2b79fd6f54f3dfe7f70be0e971b", size = 1383782, upload-time = "2026-06-25T15:04:00.29Z" }, + { url = "https://files.pythonhosted.org/packages/b8/bd/27a70984ee0a14013d046108d971327632bd5a65f0b092c66b55314dab03/fdb5lib-5.21.4.21-cp313-cp313-macosx_13_0_x86_64.whl", hash = "sha256:0a1e4349c6c1e89cd8dfbd0c87964a493e2906e25f17c967cbf3e2468b44c4a8", size = 1413311, upload-time = "2026-06-25T14:40:14.193Z" }, + { url = "https://files.pythonhosted.org/packages/b0/95/2d2a435d1476c346e66950f11b4c181bca2e7ce06fbfa2ec74f3049e01af/fdb5lib-5.21.4.21-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:7fc98326abfa122103d760138e0baaa63fd509ca23d5bb266701cc722d7e9a09", size = 1687866, upload-time = "2026-06-25T14:33:09.171Z" }, + { url = "https://files.pythonhosted.org/packages/87/9a/f5645fe4c7d7e141a3f107cdd1c946f2ee8a120b7c861527f8b41ae5ceab/fdb5lib-5.21.4.21-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:f43fd9d0aa1961905a0c973471ed98c7f33867647e5178816a36d8068fab5029", size = 1667421, upload-time = "2026-06-25T14:35:40.179Z" }, + { url = "https://files.pythonhosted.org/packages/37/80/fa2b8dde32cdf73b23deaa5c7716b9fa24e652a7664903ac4a25f9cd4451/fdb5lib-5.21.4.21-cp314-cp314-macosx_13_0_arm64.whl", hash = "sha256:59715fdb92026f6ff6aa9e450da2d61fbdbd43e4b273fa92993580a9928fa43f", size = 1383781, upload-time = "2026-06-25T14:35:24.469Z" }, + { url = "https://files.pythonhosted.org/packages/c4/b1/be84089e4e930367e8ef9ef122ba483bc5ce08f486447bc2b161ad0cb8f6/fdb5lib-5.21.4.21-cp314-cp314-macosx_13_0_x86_64.whl", hash = "sha256:792f1c9e42e2ff9133c2a495b601c412c9fd9abf7afee6ca343b0b97f41b4e0b", size = 1413312, upload-time = "2026-06-25T14:59:19.479Z" }, + { url = "https://files.pythonhosted.org/packages/48/5c/e9858d63d3d69d8539cedb3f9dfaff0eda4c7a972271eb2a4c0be39f1b5c/fdb5lib-5.21.4.21-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:ac5f25d58b80c674b3b4186a06709898bb06d05195292498820d861b1237c872", size = 1687866, upload-time = "2026-06-25T14:34:53.662Z" }, + { url = "https://files.pythonhosted.org/packages/b9/20/5c5be93b09a59d5eb546fd67f723bd6f3585784d7bd081099cb2ea864a7c/fdb5lib-5.21.4.21-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:e649d51bfc7797f3989481932b7889a72b5255d9b46722ba83505a035ba5c764", size = 1667418, upload-time = "2026-06-25T14:35:26.413Z" }, +] + [[package]] name = "filelock" version = "3.29.0" @@ -2463,6 +2493,33 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b3/38/89ba8ad64ae25be8de66a6d463314cf1eb366222074cfda9ee839c56a4b4/mdurl-0.1.2-py3-none-any.whl", hash = "sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8", size = 9979, upload-time = "2022-08-14T12:40:09.779Z" }, ] +[[package]] +name = "metkitlib" +version = "1.18.3.21" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "eccodeslib" }, + { name = "eckitlib" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/83/39/22796962d3fa1f9c1933ec435cae770a6e9c12cd84d2b60c36049ce73a6c/metkitlib-1.18.3.21-cp311-cp311-macosx_13_0_arm64.whl", hash = "sha256:2744b70ba96c0e575473147a7629d3bda32bcfb3680ee09fe80d0e874328b766", size = 948486, upload-time = "2026-06-25T14:50:14.901Z" }, + { url = "https://files.pythonhosted.org/packages/be/5f/788970e80ac0460c0ea0567c49df546d2b46f2c816be5f4f7d3b49d2d836/metkitlib-1.18.3.21-cp311-cp311-macosx_13_0_x86_64.whl", hash = "sha256:1204ce430a0a103a7a0b54ede4320643736ec9eaefb4f4a3a3184f799aed4bbd", size = 993570, upload-time = "2026-06-25T14:39:18.166Z" }, + { url = "https://files.pythonhosted.org/packages/00/75/f00cc2c7b36471c396acbaf70dfc424d94ec6e3f7ee22632d8f82714b00c/metkitlib-1.18.3.21-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:7898fca26401fe091872a3214cbdd46731b5304c3dd7b08f78c635d0c74b0f6c", size = 1038267, upload-time = "2026-06-25T14:35:14.733Z" }, + { url = "https://files.pythonhosted.org/packages/25/ce/a528262530c8c39aa8b7cbe35b6abfbc8db32dc3e75cf2aa86d000a6dde6/metkitlib-1.18.3.21-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:f2047b67f95a498d315a00a6e24cade20cc20479f7c3ca4650ad30a083d78c83", size = 1073555, upload-time = "2026-06-25T14:35:27.688Z" }, + { url = "https://files.pythonhosted.org/packages/e7/a2/775eecc029b413f019844bd8fdc6c268fdad99958e7e4246091bc2b26d2a/metkitlib-1.18.3.21-cp312-cp312-macosx_13_0_arm64.whl", hash = "sha256:6daf0c0695b3d6ee84508059e9112c7b5befe5376b55468fd3ccb87bdc0f2106", size = 948484, upload-time = "2026-06-25T15:08:01.746Z" }, + { url = "https://files.pythonhosted.org/packages/f7/98/2c28bdbc4300022418c5b05e72208434bcb608b852f1fa151d0d2cf5dbd3/metkitlib-1.18.3.21-cp312-cp312-macosx_13_0_x86_64.whl", hash = "sha256:8e1309b9dfd02b9611628f38695b3991c7c05a678566da3baab828878d004f66", size = 993570, upload-time = "2026-06-25T15:20:05.454Z" }, + { url = "https://files.pythonhosted.org/packages/22/4a/91df7945c9673e9db26c0a9fd04085ccb6c0b6287213ad9fabbda3d4241a/metkitlib-1.18.3.21-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:f93c8046368d49534844c2aaa7b45a8810ea14f291ea9241e345ddecbda22463", size = 1038267, upload-time = "2026-06-25T14:36:14.432Z" }, + { url = "https://files.pythonhosted.org/packages/22/2c/6b65b295e689ef40605c0da4e58f6042cc56061132cacefb9ba668da144d/metkitlib-1.18.3.21-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:42560798e800dec6f251c8d947d4560422dbb33cc139436954cbee674ebf8b18", size = 1073556, upload-time = "2026-06-25T14:35:54.441Z" }, + { url = "https://files.pythonhosted.org/packages/50/3a/aa7cba991330a143334f3a25958790df045a425e8f62fb6608811216b1b4/metkitlib-1.18.3.21-cp313-cp313-macosx_13_0_arm64.whl", hash = "sha256:ed70c936f9ef1ec898c7d2d4450548da9c51a289bd15ee5abbab4ac1a5b84436", size = 948486, upload-time = "2026-06-25T15:04:02.689Z" }, + { url = "https://files.pythonhosted.org/packages/b0/4e/7c28c315d4be71270120631af02b87beca7674f3b1e4e9be404a7db62a11/metkitlib-1.18.3.21-cp313-cp313-macosx_13_0_x86_64.whl", hash = "sha256:308e3e2084da51650b09657731b911f18e4edef26f16548c50c2125b61666bba", size = 993575, upload-time = "2026-06-25T14:40:17.046Z" }, + { url = "https://files.pythonhosted.org/packages/cc/67/095a563914a43250430f3c69fd235d0a0d5e8bc405e3ce2c97f2874e6a83/metkitlib-1.18.3.21-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:0155b398de2bf16759495718d9ebf43463865c2b1444f0c837ce158d67111c4c", size = 1038265, upload-time = "2026-06-25T14:33:11.495Z" }, + { url = "https://files.pythonhosted.org/packages/1f/41/e65cb9cb3949440a5b691adc2035f8897863b802ee747bddf49a2e8eb811/metkitlib-1.18.3.21-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:1d47fcd78529e67a72fad30ee7df85f2a4455245ae58cd9c38a2d3b3902311d0", size = 1073557, upload-time = "2026-06-25T14:35:43.43Z" }, + { url = "https://files.pythonhosted.org/packages/8d/b4/d99fb6f49937958ce1f2bfaa98c407b7d38e2be7dacbde51a0173cdd2504/metkitlib-1.18.3.21-cp314-cp314-macosx_13_0_arm64.whl", hash = "sha256:b2a793f597766ec1f1dfc510706ec6c7c061b5150c0f2d9f0f9914d28ab3442b", size = 948485, upload-time = "2026-06-25T14:35:28.261Z" }, + { url = "https://files.pythonhosted.org/packages/34/47/ec16aafb2e3ceea1725a5aaf33362b378c2e327c30fe96c8347ff1d94b61/metkitlib-1.18.3.21-cp314-cp314-macosx_13_0_x86_64.whl", hash = "sha256:a41b4d7a723237f2b76adc566f589d4f6f8ddaa42b8d83aaa513fdfd405f0557", size = 993573, upload-time = "2026-06-25T14:59:22.827Z" }, + { url = "https://files.pythonhosted.org/packages/4d/d0/b19d70b208731dbb88be918c066f369e5eda6d4f056e9f1dd212b134264f/metkitlib-1.18.3.21-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:13873e85ee6a0c561e5c1d0421bc80540f0ef06be3b710f53a6be6019c84397e", size = 1038265, upload-time = "2026-06-25T14:34:56.306Z" }, + { url = "https://files.pythonhosted.org/packages/fd/78/c5235fe5cd54fad3e96aa2b9e5f616dd121fbfa68283c2744d208a080606/metkitlib-1.18.3.21-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:a20603e0c9726ad5fa4f5f2a5633c6cdf6d3651d26807d8cfa812aa904c5d001", size = 1073557, upload-time = "2026-06-25T14:35:30.473Z" }, +] + [[package]] name = "mir-python" version = "1.29.0.21" @@ -3638,6 +3695,34 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/4b/2d/69abac8f838090bbecd5df894befb2c2619e7996a98ddb949db9f3b93225/pydantic_core-2.46.4-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:d51026d73fcfd93610abc7b27789c26b313920fcfb20e27462d74a7f8b06e983", size = 2193071, upload-time = "2026-05-06T13:38:08.682Z" }, ] +[[package]] +name = "pyfdb" +version = "5.21.4.21" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "fdb5lib" }, + { name = "findlibs" }, + { name = "pyyaml" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/fb/da/8a15cef3fd46a36621b51916b17b57bdbd94464b66efd93ef9118afa57b2/pyfdb-5.21.4.21-cp311-cp311-macosx_13_0_arm64.whl", hash = "sha256:13ac8d8d550d13d3780616beef413d8508f345e20c40f9d782be193be0b139d5", size = 149245, upload-time = "2026-06-25T14:50:24.918Z" }, + { url = "https://files.pythonhosted.org/packages/9e/d0/878a420309deb9f5fe3fce79322cd2f83f203715d3adce8ac95c774de949/pyfdb-5.21.4.21-cp311-cp311-macosx_13_0_x86_64.whl", hash = "sha256:2c4f12a493153771df25ea36b898f0785c23af8c5a80c6551ba2df0c4d46a996", size = 156545, upload-time = "2026-06-25T14:39:28.737Z" }, + { url = "https://files.pythonhosted.org/packages/ab/7d/f3a912f9cf22b4aab9828d720346c81635e304d925cc35c852a508f86297/pyfdb-5.21.4.21-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:de41c16f3c0a3b4c4d517fe48597a145b475df413bfcb23ce121d19bf09fc6e4", size = 143177, upload-time = "2026-06-25T14:35:23.566Z" }, + { url = "https://files.pythonhosted.org/packages/f1/ef/f11938857d4947e079661d03ec2814ccfa7f0cd9cd5e0455bffe798ed450/pyfdb-5.21.4.21-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:ea0eb7f2410c8e56a6f2493de535c8d7e323bf97ec656627ba1b145f80b8e97b", size = 149461, upload-time = "2026-06-25T14:35:40.975Z" }, + { url = "https://files.pythonhosted.org/packages/5b/8f/29a20431c99615dcac06627a9a9b9560e12a9a88751fe95d5f97696a34d3/pyfdb-5.21.4.21-cp312-cp312-macosx_13_0_arm64.whl", hash = "sha256:f4c386a219bb39952e890f550eeb52d5ed2cbcd92b14fcc9102a3ccca21311d2", size = 166134, upload-time = "2026-06-25T15:08:10.709Z" }, + { url = "https://files.pythonhosted.org/packages/d7/d6/27569905374b1c927066b742a3e22f76276cfee55d5594fa1ab8a440ff05/pyfdb-5.21.4.21-cp312-cp312-macosx_13_0_x86_64.whl", hash = "sha256:111db0977abd0ac822246ff06002b504e715f3debc1ae1f103a61b639f851e0c", size = 175417, upload-time = "2026-06-25T15:20:16.483Z" }, + { url = "https://files.pythonhosted.org/packages/28/29/66a73a0d052d94e3764fd338846515fef9bc80cd7a647054ed7606712554/pyfdb-5.21.4.21-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:6fa072926285b0b97bb097c1a8317c21ef0176887ba85957245d7cfb315b379f", size = 143569, upload-time = "2026-06-25T14:36:23.404Z" }, + { url = "https://files.pythonhosted.org/packages/83/73/3159c403f5e36b39ae90848924df531ba8aae067e7ad7bab4c3d5d885984/pyfdb-5.21.4.21-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:01be2c0f5722746c689f5a905b7100dc0dbcfa5baf52fc57385b48080d248563", size = 150025, upload-time = "2026-06-25T14:36:06.596Z" }, + { url = "https://files.pythonhosted.org/packages/a4/a5/d7ede9ce50b84f38251225b7aa9d4b718738f1ec38bbee304b482a75125c/pyfdb-5.21.4.21-cp313-cp313-macosx_13_0_arm64.whl", hash = "sha256:90d0fced85646b3813cbf02c67dad3e459d8e881403c62d5507dc25a20769e2b", size = 166136, upload-time = "2026-06-25T15:04:11.847Z" }, + { url = "https://files.pythonhosted.org/packages/6e/cc/26717b54d031fc5b954e9ccae4b1a16e580cd9bef39825498381fe4ff9b9/pyfdb-5.21.4.21-cp313-cp313-macosx_13_0_x86_64.whl", hash = "sha256:17c408cc9d9f2d63ebb6dda993336b572d124a9a34e0edc83410c244a9161038", size = 175435, upload-time = "2026-06-25T14:40:28.179Z" }, + { url = "https://files.pythonhosted.org/packages/b8/82/a26a36ff544620e2f27a529d872d71b5b925083934baab228a4df1b1d790/pyfdb-5.21.4.21-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:1e0b85ec7ee5963abacda33e8c08a71415ff3d73d03ca9ab9edf6a1919b522b6", size = 143606, upload-time = "2026-06-25T14:33:20.076Z" }, + { url = "https://files.pythonhosted.org/packages/cf/0b/21be146d12fc8b4e42a62ef2f64b8d3e97c68c54998a7d53c85028358361/pyfdb-5.21.4.21-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:83d2425d33a2bacdf0b05136a374c05be0dd34f1bf34e4b4996f1d8ab6e38645", size = 149988, upload-time = "2026-06-25T14:35:55.227Z" }, + { url = "https://files.pythonhosted.org/packages/7b/27/ce41bef04b9aa5ecce4914a10cd52e0e2d2e117fa32646b833861d20674f/pyfdb-5.21.4.21-cp314-cp314-macosx_13_0_arm64.whl", hash = "sha256:7924fd9176597fc51cd74e6037632af938795c25d26fd32b9a0a8ab1673e6744", size = 165426, upload-time = "2026-06-25T14:35:38.449Z" }, + { url = "https://files.pythonhosted.org/packages/e6/42/4ee89d919b03732cdb4cfda7572794cc1766a52ab7827cf0e8a44cb9803e/pyfdb-5.21.4.21-cp314-cp314-macosx_13_0_x86_64.whl", hash = "sha256:21c19111739778c1381cba57dd0e3e03a4b3991a4f884e1887fa7552853205d2", size = 175474, upload-time = "2026-06-25T14:59:33.672Z" }, + { url = "https://files.pythonhosted.org/packages/f6/9c/936cdaad1e92753b92f7cc8acc5cb77fc470c15f4be748c0dfb0cc0228da/pyfdb-5.21.4.21-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:4fbb88afaae3f3f7152f2deebe962ae42a93baa659b341c91eade500ac2a4f9c", size = 143661, upload-time = "2026-06-25T14:35:05.139Z" }, + { url = "https://files.pythonhosted.org/packages/2c/c5/1d11a305af088097983401130c78eb611f5dacc2318e8db2971065364a53/pyfdb-5.21.4.21-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:49d4ec5137482101f207468e08cb0a7c7755485630ad682633fe04fc459edc1d", size = 150001, upload-time = "2026-06-25T14:35:43.604Z" }, +] + [[package]] name = "pygments" version = "2.20.0" diff --git a/workflow/rules/inference.smk b/workflow/rules/inference.smk index 8c3d91b2..91c047ee 100644 --- a/workflow/rules/inference.smk +++ b/workflow/rules/inference.smk @@ -127,6 +127,15 @@ rule inference_prepare_env: git@github.com:meteoswiss/eccodes-cosmo-mars.git \ "$VENV/data/share/eccodes-cosmo-mars" echo "[$(date)] Cloned eccodes-cosmo-mars definitions" + + # metkitlib (pulled in by pyfdb) ships its own language.yaml, used to + # validate fdb.retrieve() requests. Our custom `model` values (added + # purely for archival, via eccodes-cosmo-mars's marsModel concepts) + # were never registered there, so retrieve() rejects them even though + # they're valid archived keys -- fdb.list() has no such validation, so + # this only bites downstream reads, not archival or completeness checks. + echo "[$(date)] Patching metkitlib's language.yaml for custom FDB model names..." + python resources/fdb/patch_metkit_language.py "$VENV" fi echo "[$(date)] Compiling Python bytecode..." @@ -280,9 +289,14 @@ checkpoint inference_check_fdb: fi # Mount the venv (pyfdb, python) and the FDB uenv (native libs, config, # at its required /user-environment path) side by side in one call. + # squashfs-mount needs the mountpoint dir to already exist -- /user-environment + # is a pre-provisioned system directory, but there's no such slot for a second + # image, so the venv gets a job-local directory we create ourselves. + VENV_MOUNT=$(mktemp -d) + trap "rmdir '$VENV_MOUNT'" EXIT FDB_SQUASHFS=$(uenv image inspect --format='{{sqfs}}' {params.fdb_uenv}) - squashfs-mount {params.image_path}:/venv-environment "$FDB_SQUASHFS:/user-environment" -- bash -c ' - source /venv-environment/bin/activate + squashfs-mount {params.image_path}:$VENV_MOUNT "$FDB_SQUASHFS:/user-environment" -- bash -c ' + source '"$VENV_MOUNT"'/bin/activate export LD_LIBRARY_PATH=/user-environment/env/{params.fdb_view}/lib64:/user-environment/env/{params.fdb_view}/lib${{LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}} python {input.script} '"$GLOBAL_ARG"' \ --fdb-root {params.fdb_root_local} \ @@ -298,14 +312,28 @@ checkpoint inference_check_fdb: """ +def _resolve_fdb_root_for_run(run_id): + """Which FDB root this run's own inference_execute would write fresh output + into: the global one if write_to_global_fdb is set (and configured), + otherwise the local per-checkpoint root. Same logic as inference_execute's + own `fdb_root` param -- factored out so the read side can agree with it.""" + global_root, local_root = _get_fdb_roots(run_id) + env_id = RUN_CONFIGS[run_id]["env_id"] + if get_resource_by_env(env_id, "write_to_global_fdb", False) and global_root: + return global_root + return local_root + + def _get_inference_status(wc): """Return (okfile_path, fdb_root_or_None) for (run_id, init_time). - 'hit' -> data already complete in FDB; okfile is the check's status file - (inference_execute is never requested), fdb_root points at the - root holding it. - 'miss' or FDB not configured -> falls through to inference_execute's okfile, - fdb_root is None (data will be local grib / freshly written FDB). + fdb_root is non-None whenever this run's own config writes to FDB at all + (its output.tee has fdb: targets) -- whether the data was already there + (a check 'hit') or is about to be freshly written by inference_execute + (a 'miss', or FDB-check not configured but still an FDB-writing run). + It's None only when this run's config has no FDB targets at all (plain + local-grib output), in which case downstream code reads local grib as + before. """ env_id = RUN_CONFIGS[wc.run_id]["env_id"] if _get_fdb_uenv_for_env(env_id)[0]: @@ -315,7 +343,11 @@ def _get_inference_status(wc): status, _, root = Path(status_file).read_text().strip().partition("\t") if status == "hit": return str(status_file), root - return str(OUT_ROOT / f"logs/inference_execute/{wc.run_id}-{wc.init_time}.ok"), None + okfile = str(OUT_ROOT / f"logs/inference_execute/{wc.run_id}-{wc.init_time}.ok") + fdb_root = ( + _resolve_fdb_root_for_run(wc.run_id) if _count_fdb_blocks(wc.run_id) > 0 else None + ) + return okfile, fdb_root def _get_inference_okfile(wc): @@ -454,12 +486,7 @@ rule inference_execute: fdb_configured=lambda wc: bool( _get_fdb_uenv_for_env(RUN_CONFIGS[wc.run_id]["env_id"])[0] ), - fdb_root=lambda wc: ( - _get_fdb_roots(wc.run_id)[0] - if get_resource(wc, "write_to_global_fdb", False) - and _get_fdb_roots(wc.run_id)[0] - else _get_fdb_roots(wc.run_id)[1] - ), + fdb_root=lambda wc: _resolve_fdb_root_for_run(wc.run_id), fdb_schema=str(Path("resources/fdb/realtime-varda.schema").resolve()), # fmt: off shell: @@ -524,28 +551,42 @@ FDBEOF CMD_ARGS+=(runner.parallel.cluster=slurm) fi - {params.srun_prefix} srun \ - --unbuffered \ - --partition={resources.slurm_partition} \ - --cpus-per-task={resources.cpus_per_task} \ - --mem-per-cpu={resources.mem_mb_per_cpu} \ - --time={resources.runtime} \ - --gres={resources.gres} \ - --ntasks={resources.ntasks} \ - anemoi-inference run config.yaml "${{CMD_ARGS[@]}}" + anemoi-inference run config.yaml "${{CMD_ARGS[@]}}" }} export -f _run_inference - # Mount the venv and (if configured) the FDB uenv side by side in one - # call: the FDB uenv's absolute internal symlinks require it to be at - # /user-environment, so the venv takes a different mountpoint instead. + # squashfs-mount needs the mountpoint dir to already exist, and -- since srun + # may dispatch to a different physical node -- it must be on the shared + # filesystem, not /tmp (node-local). /user-environment is a pre-provisioned + # system directory present on every node for the FDB uenv's own absolute + # internal symlinks; the venv gets a job-local dir here instead. + # + # srun must be the OUTER call, dispatching a (dual) squashfs-mount as its + # task, not the other way around: wrapping srun inside a dual-image + # squashfs-mount breaks Slurm's plugin initialization on this system + # ("Plugin initialization failed") -- verified empirically. A single-image + # mount around srun is fine; two images at once are not. + VENV_MOUNT="{params.workdir}/.venv-mount" + mkdir -p "$VENV_MOUNT" + trap "rmdir '$VENV_MOUNT'" EXIT + + MOUNT_ARGS=("{params.env_path}:$VENV_MOUNT") + INNER_CMD="_run_inference $VENV_MOUNT" if [ "{params.fdb_configured}" = "True" ]; then FDB_SQUASHFS=$(uenv image inspect --format='{{sqfs}}' {params.fdb_uenv}) - squashfs-mount {params.env_path}:/venv-environment "$FDB_SQUASHFS:/user-environment" -- bash -c '_run_inference /venv-environment /user-environment' - else - squashfs-mount {params.env_path}:/venv-environment -- bash -c '_run_inference /venv-environment' + MOUNT_ARGS+=("$FDB_SQUASHFS:/user-environment") + INNER_CMD="$INNER_CMD /user-environment" fi - ) >{log} 2>&1 - touch {output.okfile} + + {params.srun_prefix} srun \ + --unbuffered \ + --partition={resources.slurm_partition} \ + --cpus-per-task={resources.cpus_per_task} \ + --mem-per-cpu={resources.mem_mb_per_cpu} \ + --time={resources.runtime} \ + --gres={resources.gres} \ + --ntasks={resources.ntasks} \ + squashfs-mount "${{MOUNT_ARGS[@]}}" -- bash -c "$INNER_CMD" + ) >{log} 2>&1 && touch {output.okfile} """ # fmt: on diff --git a/workflow/rules/plot.smk b/workflow/rules/plot.smk index fe43605b..10ef02ec 100644 --- a/workflow/rules/plot.smk +++ b/workflow/rules/plot.smk @@ -64,17 +64,46 @@ rule plot_meteogram: if config.get("lapse_rate_correction", True) else "" ), + fdb_configured=lambda wc: bool( + _get_fdb_uenv_for_env(RUN_CONFIGS[wc.run_id]["env_id"])[0] + ), + fdb_forecast=lambda wc: ( + "fdb:" + + ( + _get_fdb_roots(wc.run_id)[0] + if get_resource(wc, "write_to_global_fdb", False) + and _get_fdb_roots(wc.run_id)[0] + else _get_fdb_roots(wc.run_id)[1] + ) + if _get_fdb_uenv_for_env(RUN_CONFIGS[wc.run_id]["env_id"])[0] + else "" + ), shell: """ set -euo pipefail - export ECCODES_DEFINITION_PATH=$(realpath .venv/share/eccodes-cosmo-resources/definitions) BASELINE_ROOTS=({params.baseline_roots:q}) BASELINE_STEPS=({params.baseline_steps:q}) BASELINE_LABELS=({params.baseline_labels:q}) + # If this run's environment archives to FDB, its forecast output lives + # there (never in local grib -- see workflow/rules/inference.smk), so + # always prefer the FDB source over the (nonexistent) local grib dir. + # Reading from FDB needs pyfdb (an evalml dependency, see pyproject.toml) + # and eccodes-cosmo-mars's MARS-namespace GRIB concepts (a repo-level + # clone, resources/fdb/patch_metkit_language.py's docstring has details) + # -- no squashfs-mount of the per-checkpoint venv/FDB uenv needed, unlike + # inference_execute: pyfdb bundles its own native FDB libraries. + FORECAST_SRC={params.fcst_grib:q} + if [ "{params.fdb_configured}" = "True" ]; then + FORECAST_SRC={params.fdb_forecast:q} + export ECCODES_DEFINITION_PATH="$(realpath eccodes-cosmo-mars/definitions):$(realpath .venv/share/eccodes-cosmo-resources/definitions)" + else + export ECCODES_DEFINITION_PATH=$(realpath .venv/share/eccodes-cosmo-resources/definitions) + fi + CMD_ARGS=( - --forecast {params.fcst_grib:q} + --forecast "$FORECAST_SRC" --forecast_steps {params.fcst_steps:q} --forecast_label {params.fcst_label:q} --analysis {params.truth_root:q} @@ -124,15 +153,44 @@ rule plot_forecast_frame: (Path(OUT_ROOT) / f"data/runs/{wc.run_id}/{wc.init_time}/frames").resolve() ), accu=lambda wc: int(RUN_CONFIGS[wc.run_id]["steps"].split("/")[2]), + fdb_configured=lambda wc: bool( + _get_fdb_uenv_for_env(RUN_CONFIGS[wc.run_id]["env_id"])[0] + ), + fdb_forecast=lambda wc: ( + "fdb:" + + ( + _get_fdb_roots(wc.run_id)[0] + if get_resource(wc, "write_to_global_fdb", False) + and _get_fdb_roots(wc.run_id)[0] + else _get_fdb_roots(wc.run_id)[1] + ) + if _get_fdb_uenv_for_env(RUN_CONFIGS[wc.run_id]["env_id"])[0] + else "" + ), shell: """ - export ECCODES_DEFINITION_PATH=$(realpath .venv/share/eccodes-cosmo-resources/definitions) - python {input.script} \ - --input {params.grib_out_dir:q} --date {wildcards.init_time:q} \ - --param {wildcards.param:q} --leadtime {wildcards.leadtime:q} \ - --regions_json {params.regions_json:q} \ - --outdir {params.outdir:q} \ - --accu {params.accu} >{log} 2>&1 + set -euo pipefail + + # See plot_meteogram for why this always wins over the (nonexistent) + # local grib dir when the run's environment archives to FDB, and why no + # squashfs-mount is needed to read it (pyfdb bundles its own native libs). + INPUT_SRC={params.grib_out_dir:q} + if [ "{params.fdb_configured}" = "True" ]; then + INPUT_SRC={params.fdb_forecast:q} + export ECCODES_DEFINITION_PATH="$(realpath eccodes-cosmo-mars/definitions):$(realpath .venv/share/eccodes-cosmo-resources/definitions)" + else + export ECCODES_DEFINITION_PATH=$(realpath .venv/share/eccodes-cosmo-resources/definitions) + fi + + CMD_ARGS=( + --input "$INPUT_SRC" --date {wildcards.init_time:q} + --param {wildcards.param:q} --leadtime {wildcards.leadtime:q} + --regions_json {params.regions_json:q} + --outdir {params.outdir:q} + --accu {params.accu} + ) + + python {input.script} "${{CMD_ARGS[@]}}" >{log} 2>&1 """ diff --git a/workflow/rules/verification.smk b/workflow/rules/verification.smk index d0310f68..43e267de 100644 --- a/workflow/rules/verification.smk +++ b/workflow/rules/verification.smk @@ -97,21 +97,51 @@ rule verification_metrics: if config.get("lapse_rate_correction", True) else "" ), + fdb_configured=lambda wc: bool( + _get_fdb_uenv_for_env(RUN_CONFIGS[wc.run_id]["env_id"])[0] + ), + fdb_forecast=lambda wc: ( + "fdb:" + + ( + _get_fdb_roots(wc.run_id)[0] + if get_resource(wc, "write_to_global_fdb", False) + and _get_fdb_roots(wc.run_id)[0] + else _get_fdb_roots(wc.run_id)[1] + ) + if _get_fdb_uenv_for_env(RUN_CONFIGS[wc.run_id]["env_id"])[0] + else "" + ), shell: """ - export ECCODES_DEFINITION_PATH=$(realpath .venv/share/eccodes-cosmo-resources/definitions) - uv run {input.script} \ - --forecast {params.grib_out_dir} \ - --truth {params.truth} \ - --reftime {wildcards.init_time} \ - --steps "{params.fcst_steps}" \ - --source_id "{wildcards.run_id}" \ - --truth_source_id "{params.truth_source_id}" \ - --regions '{params.regions}' \ - --params "{params.experiment_params}" \ - --threshold_dict "{params.threshold_dict}" \ - {params.lapse_rate_flag} \ - --output {output} >{log} 2>&1 + set -euo pipefail + + # See plot_meteogram (plot.smk) for why this always wins over the + # (nonexistent) local grib dir when the run's environment archives to + # FDB, and why no squashfs-mount is needed to read it (pyfdb bundles + # its own native libs). + FORECAST_SRC={params.grib_out_dir:q} + if [ "{params.fdb_configured}" = "True" ]; then + FORECAST_SRC={params.fdb_forecast:q} + export ECCODES_DEFINITION_PATH="$(realpath eccodes-cosmo-mars/definitions):$(realpath .venv/share/eccodes-cosmo-resources/definitions)" + else + export ECCODES_DEFINITION_PATH=$(realpath .venv/share/eccodes-cosmo-resources/definitions) + fi + + CMD_ARGS=( + --forecast "$FORECAST_SRC" + --truth {params.truth:q} + --reftime {wildcards.init_time:q} + --steps {params.fcst_steps:q} + --source_id {wildcards.run_id:q} + --truth_source_id {params.truth_source_id:q} + --regions {params.regions:q} + --params {params.experiment_params:q} + --threshold_dict {params.threshold_dict:q} + {params.lapse_rate_flag} + --output {output:q} + ) + + uv run python {input.script} "${{CMD_ARGS[@]}}" >{log} 2>&1 """ @@ -222,19 +252,54 @@ rule verification_scoremaps: fcst_label=lambda wc: RUN_CONFIGS[wc.run_id].get("label"), fcst_steps=lambda wc: RUN_CONFIGS[wc.run_id]["steps"], truth_label=(config.get("truth") or {}).get("label", ""), - reftimes=" ".join(t.strftime("%Y%m%d%H%M") for t in REFTIMES), + reftimes=[t.strftime("%Y%m%d%H%M") for t in REFTIMES], run_root=lambda wc: (Path(OUT_ROOT) / f"data/runs/{wc.run_id}").resolve(), + fdb_configured=lambda wc: bool( + _get_fdb_uenv_for_env(RUN_CONFIGS[wc.run_id]["env_id"])[0] + ), + # Bare root (no "fdb:" marker prefix): --fdb_root is consumed by + # iter_init_dirs/fdb_has_data, which add the marker themselves when + # constructing load_forecast_data's root argument for a discovered + # FDB-resident reftime -- unlike plot.smk/verification_metrics' single + # --forecast argument, which IS load_forecast_data's root directly. + fdb_root=lambda wc: ( + _get_fdb_roots(wc.run_id)[0] + if get_resource(wc, "write_to_global_fdb", False) + and _get_fdb_roots(wc.run_id)[0] + else _get_fdb_roots(wc.run_id)[1] + ) + if _get_fdb_uenv_for_env(RUN_CONFIGS[wc.run_id]["env_id"])[0] + else "", shell: """ - export ECCODES_DEFINITION_PATH=$(realpath .venv/share/eccodes-cosmo-resources/definitions) - uv run {input.script} \ - --run_root {params.run_root} \ - --reftimes {params.reftimes} \ - --truth {input.truth} \ - --step {wildcards.leadtime} \ - --steps "{params.fcst_steps}" \ - --param {wildcards.param} \ - --output {output} >{log} 2>&1 + set -euo pipefail + + # See plot_meteogram (plot.smk) for why no squashfs-mount is needed to + # read from FDB (pyfdb bundles its own native libs). + if [ "{params.fdb_configured}" = "True" ]; then + export ECCODES_DEFINITION_PATH="$(realpath eccodes-cosmo-mars/definitions):$(realpath .venv/share/eccodes-cosmo-resources/definitions)" + else + export ECCODES_DEFINITION_PATH=$(realpath .venv/share/eccodes-cosmo-resources/definitions) + fi + + # --reftimes is nargs="+" in argparse: pass each timestamp as its own + # array element (not one space-joined string) so it fans out correctly. + REFTIMES_ARR=({params.reftimes:q}) + + CMD_ARGS=( + --run_root {params.run_root:q} + --reftimes "${{REFTIMES_ARR[@]}}" + --truth {input.truth:q} + --step {wildcards.leadtime:q} + --steps {params.fcst_steps:q} + --param {wildcards.param:q} + --output {output:q} + ) + if [ -n "{params.fdb_root}" ]; then + CMD_ARGS+=(--fdb_root {params.fdb_root:q}) + fi + + uv run python {input.script} "${{CMD_ARGS[@]}}" >{log} 2>&1 """ diff --git a/workflow/scripts/plot_forecast_frame.py b/workflow/scripts/plot_forecast_frame.py index ee8e564e..472611f5 100644 --- a/workflow/scripts/plot_forecast_frame.py +++ b/workflow/scripts/plot_forecast_frame.py @@ -1,6 +1,7 @@ import json import logging from argparse import ArgumentParser +from datetime import datetime from pathlib import Path import cartopy.crs as ccrs @@ -14,7 +15,7 @@ from plotting import get_projection from plotting import StatePlotter from plotting.colormap_defaults import CMAP_DEFAULTS -from plotting.compat import load_state_from_grib +from plotting.compat import load_state_from_grib, load_state_from_fdb LOG = logging.getLogger(__name__) logging.basicConfig( @@ -119,7 +120,7 @@ def main(): ) args = parser.parse_args() - grib_dir = Path(args.input) + is_fdb = args.input.startswith("fdb:") init_time = args.date lead_time = args.leadtime param = args.param @@ -143,43 +144,72 @@ def main(): else: paramlist = [param] - # Load grib once — shared across all region plots - # TODO: fix file pattern & globbing - grib_file = list(grib_dir.glob(f"2*_{lead_time}.grib")) - if not grib_file: - grib_file = list(grib_dir.glob(f"2*_{lead_time:03}.grib")) - if not grib_file: - LOG.warning( - "No GRIB file found for lead_time=%s in %s (model may not write initial state)." - " Creating empty placeholder frames.", - lead_time, - grib_dir, - ) + def _empty_frames(): for region_name in regions: outfn = outdir / f"frame_{lead_time}_{param}_{region_name}.png" outfn.touch() - return - grib_file = Path(grib_file[0]) - LOG.info("Loading grib file %s", grib_file) - state = load_state_from_grib(grib_file, paramlist=paramlist) + + # Load state once — shared across all region plots + if is_fdb: + fdb_root = args.input[len("fdb:") :] + reftime = datetime.strptime(init_time, "%Y%m%d%H%M") + LOG.info("Loading forecast state from FDB %s", fdb_root) + try: + state = load_state_from_fdb(fdb_root, reftime, lead_time, paramlist=paramlist) + except FileNotFoundError: + LOG.warning( + "No FDB data for lead_time=%s in %s (model may not write initial state)." + " Creating empty placeholder frames.", + lead_time, + fdb_root, + ) + _empty_frames() + return + else: + grib_dir = Path(args.input) + # TODO: fix file pattern & globbing + grib_file = list(grib_dir.glob(f"2*_{lead_time}.grib")) + if not grib_file: + grib_file = list(grib_dir.glob(f"2*_{lead_time:03}.grib")) + if not grib_file: + LOG.warning( + "No GRIB file found for lead_time=%s in %s (model may not write initial state)." + " Creating empty placeholder frames.", + lead_time, + grib_dir, + ) + _empty_frames() + return + grib_file = Path(grib_file[0]) + LOG.info("Loading grib file %s", grib_file) + state = load_state_from_grib(grib_file, paramlist=paramlist) # tp is accumulated from start of forecast; de-accumulate to get period [lt-accu, lt] if param == "TOT_PREC": prev_lt = lead_time - accu if prev_lt > 0: - prev_grib_files = list(grib_dir.glob(f"2*_{prev_lt}.grib")) - if not prev_grib_files: - prev_grib_files = list(grib_dir.glob(f"2*_{prev_lt:03d}.grib")) - prev_grib_file = Path(prev_grib_files[0]) - LOG.info( - "De-accumulating TOT_PREC: loading previous grib file %s", - prev_grib_file, - ) - prev_state = load_state_from_grib(prev_grib_file, paramlist=paramlist) - state["fields"]["TOT_PREC"] = ( - state["fields"]["TOT_PREC"] - - prev_state["fields"]["TOT_PREC"][: len(state["fields"]["TOT_PREC"])] - ) + if is_fdb: + try: + prev_state = load_state_from_fdb( + fdb_root, reftime, prev_lt, paramlist=paramlist + ) + except FileNotFoundError: + prev_state = None + else: + prev_grib_files = list(grib_dir.glob(f"2*_{prev_lt}.grib")) + if not prev_grib_files: + prev_grib_files = list(grib_dir.glob(f"2*_{prev_lt:03d}.grib")) + prev_state = ( + load_state_from_grib(Path(prev_grib_files[0]), paramlist=paramlist) + if prev_grib_files + else None + ) + if prev_state is not None: + LOG.info("De-accumulating TOT_PREC using previous lead_time=%s", prev_lt) + state["fields"]["TOT_PREC"] = ( + state["fields"]["TOT_PREC"] + - prev_state["fields"]["TOT_PREC"][: len(state["fields"]["TOT_PREC"])] + ) # Preprocess field once — shared across all region plots field, units_override = preprocess_field(param, state) diff --git a/workflow/scripts/verification_scoremaps.py b/workflow/scripts/verification_scoremaps.py index d640352a..e17febcc 100644 --- a/workflow/scripts/verification_scoremaps.py +++ b/workflow/scripts/verification_scoremaps.py @@ -39,6 +39,7 @@ import xarray as xr from data_input import ( + fdb_has_data, load_forecast_data, load_truth_data, open_truth_zarr, @@ -75,12 +76,24 @@ def _season_of(dt: datetime) -> str: # --------------------------------------------------------------------------- -def iter_init_dirs(run_root: Path) -> list[tuple[datetime, Path]]: +def iter_init_dirs( + run_root: Path, fdb_root: str | None = None +) -> list[tuple[datetime, Path | str]]: """Return ``(reftime, grib_dir)`` pairs for every complete init time. Expects subdirectories named ``YYYYMMDDHHMI`` directly under *run_root*. GRIB files may live either directly in the init-time directory or inside a ``grib/`` subdirectory. + + If `fdb_root` is given, init-time directories with no local GRIB output + (e.g. runs that archive directly to FDB and never write local grib/ at + all) are additionally checked against it and returned as an ``"fdb:"`` + marker string in place of a real path -- see data_input.load_forecast_data's + routing on that prefix, which every caller of this function's output + already goes through. Presence is checked with fdb_has_data(), a bare + existence check: completeness was already established upstream, by the + Snakemake dependency on this reftime's inference_check_fdb/inference_execute + okfile, before this script ever runs. """ result = [] for d in sorted(run_root.iterdir()): @@ -91,10 +104,12 @@ def iter_init_dirs(run_root: Path) -> list[tuple[datetime, Path]]: except ValueError: continue grib_dir = d / "grib" if (d / "grib").is_dir() else d - if not any(grib_dir.glob("*.grib")): + if any(grib_dir.glob("*.grib")): + result.append((reftime, grib_dir)) + elif fdb_root is not None and fdb_has_data(fdb_root, reftime): + result.append((reftime, f"fdb:{fdb_root}")) + else: LOG.debug("No GRIB files in %s, skipping", grib_dir) - continue - result.append((reftime, grib_dir)) return result @@ -139,7 +154,7 @@ def main(args: Namespace) -> None: ] LOG.info("Using %d baseline init times from --reftimes", len(init_items)) else: - init_items = iter_init_dirs(args.run_root) + init_items = iter_init_dirs(args.run_root, fdb_root=args.fdb_root) LOG.info("Found %d init time directories", len(init_items)) # Restrict to the experiment's configured init times, and require that @@ -454,6 +469,17 @@ def _strat_da(compute_fn) -> xr.DataArray: "--reftimes." ), ) + parser.add_argument( + "--fdb_root", + type=str, + default=None, + help=( + "FDB root holding this run's output, if any is configured for its " + "environment. Init times with no local GRIB output under --run_root " + "are additionally checked against it (see iter_init_dirs). Ignored " + "for --baseline_root." + ), + ) parser.add_argument( "--member", type=str, From 53e37f5cad20a67681315353630783bcdbc4d35e Mon Sep 17 00:00:00 2001 From: Carlos Osuna Date: Tue, 15 Sep 2026 13:23:42 +0200 Subject: [PATCH 04/21] add a language patch --- resources/fdb/patch_metkit_language.py | 71 ++++++++++++++++++++++++++ 1 file changed, 71 insertions(+) create mode 100644 resources/fdb/patch_metkit_language.py diff --git a/resources/fdb/patch_metkit_language.py b/resources/fdb/patch_metkit_language.py new file mode 100644 index 00000000..aa3404f0 --- /dev/null +++ b/resources/fdb/patch_metkit_language.py @@ -0,0 +1,71 @@ +"""Patch metkitlib's bundled language.yaml to recognize VARDA-SINGLE's custom +`model` values for FDB *retrieve* requests. + +`model` is not a standard MARS/metkit key -- it was added purely for FDB archival, +via the `marsModel` concepts in eccodes-cosmo-mars's local.215.def/local.98.def +(see resources/fdb/realtime-varda.schema and the VARDA-SINGLE FDB archival notes). +`fdb.list()` doesn't validate key values against any vocabulary, so archival and +completeness checks work fine without this patch. But `fdb.retrieve()` parses +requests through metkit's MarsLanguage, which validates enum-typed keys like +`model` against language.yaml's fixed value list -- and rejects anything not +already registered there, even though the value is a perfectly valid archived key. + +Run once per venv build, after `metkitlib` is installed (it ships its own copy of +language.yaml under its package data; there's no config-based override for this, +unlike ECCODES_DEFINITION_PATH/FDB5_CONFIG_FILE, so the file has to be edited in place). +""" + +import argparse +import re +import sys +from pathlib import Path + +# Custom VARDA-SINGLE model names, per marsModel concepts in eccodes-cosmo-mars's +# local.215.def (ICON, GPI=180) / local.98.def (IFS, GPI=180). "varda-single-g" is +# the global-domain variant seen from the temporal downscaler's IFS-grid output. +_CUSTOM_MODELS = ["varda-single", "varda-single-g"] + +# Structural anchor, not a specific value: metkitlib/share/metkit/language.yaml's +# top-level ` model:` key ends with the default (no-context) `type: enum` block, +# listing recognized model names -- its last entry (and thus the exact line to +# anchor on) shifts between metkitlib releases as upstream adds new models, so +# this matches the *end of the model: block* (a blank line before the next +# top-level ` key:`) instead of any one entry. +_BLOCK_END_RE = re.compile(r"(\n model:\n(?:.*\n)*?)(\n [a-zA-Z_]+:)") + + +def patch(language_yaml: Path) -> None: + text = language_yaml.read_text() + if all(f"- [{m}]" in text for m in _CUSTOM_MODELS): + print(f"[patch_metkit_language] {language_yaml} already patched, skipping") + return + m = _BLOCK_END_RE.search(text) + if m is None: + raise RuntimeError( + f"Could not locate the end of the 'model:' block in {language_yaml} -- " + "metkitlib's language.yaml layout may have changed; update " + "_BLOCK_END_RE in this script." + ) + addition = "\n".join(f" - [{model}]" for model in _CUSTOM_MODELS) + text = text[: m.end(1)] + addition + text[m.end(1) :] + language_yaml.write_text(text) + print(f"[patch_metkit_language] added {_CUSTOM_MODELS} to {language_yaml}") + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "venv", type=Path, help="Path to the venv root (contains lib/python*/site-packages)" + ) + args = parser.parse_args() + + matches = list(args.venv.glob("lib/python*/site-packages/metkitlib/share/metkit/language.yaml")) + if not matches: + raise RuntimeError(f"metkitlib language.yaml not found under {args.venv}") + for language_yaml in matches: + patch(language_yaml) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) From 1c2b228a9ba1c1e3fe789f911426f1c7db305ff1 Mon Sep 17 00:00:00 2001 From: Carlos Osuna Date: Fri, 16 Oct 2026 14:59:19 +0200 Subject: [PATCH 05/21] chore: ignore core dumps and local eccodes-cosmo-mars checkout Neither belongs in version control: core dumps are debug artifacts from crash investigation, and the eccodes-cosmo-mars checkout (needed locally by plot.smk/verification.smk for FDB reads outside the squashfs-mounted venv) is a clone of an external repo, not vendored content. --- .gitignore | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/.gitignore b/.gitignore index 50b2ff7d..61983782 100644 --- a/.gitignore +++ b/.gitignore @@ -55,3 +55,12 @@ uv.lock # evalml .evalml_snakemake_cmd.txt + +# Core dumps +core_* + +# Local clone of MeteoSwiss's eccodes-cosmo-mars (MARS-namespace GRIB +# definitions needed for FDB reads outside the squashfs-mounted venv, see +# workflow/rules/plot.smk and verification.smk). Not vendored: each checkout +# clones it locally, same as the one baked into the inference venv. +/eccodes-cosmo-mars/ From 7a577103401d90507320898630297609ca54f081 Mon Sep 17 00:00:00 2001 From: Carlos Osuna Date: Fri, 16 Oct 2026 14:59:28 +0200 Subject: [PATCH 06/21] fix: rewrite FDB read path to request exact params via earthkit-data The previous implementation used pyfdb's list() to discover every archived parameter and then fetched each one individually, regardless of which params the caller actually asked for; bytes were also round-tripped through a temp file. Rewrite to issue a single request per (levtype, timespan) group for exactly the requested params, via earthkit-data's "fdb" source (ekd.from_source) rather than raw pyfdb, with no intermediate file. Also fixes FieldList.group_by("typeOfLevel") returning one group regardless of distinct values (manual grouping instead), and adds the timespan=fs key required for TOT_PREC (otherwise silently empty for any requested step). --- src/data_input/__init__.py | 258 ++++++++++++++++++++++++++----------- src/plotting/compat.py | 28 ++-- 2 files changed, 200 insertions(+), 86 deletions(-) diff --git a/src/data_input/__init__.py b/src/data_input/__init__.py index cef9064e..ee531e31 100644 --- a/src/data_input/__init__.py +++ b/src/data_input/__init__.py @@ -1,8 +1,6 @@ -import io import json import logging import re -import tempfile import urllib.request from datetime import datetime, timedelta from pathlib import Path @@ -59,6 +57,71 @@ _FDB_MODEL = "varda-single" _FDB_GLOBAL_MODEL = "varda-single-g" +# Numeric FDB "param" codes for VARDA-SINGLE's custom MARS namespace (GPI=180 +# local table, eccodes-cosmo-mars's local.215.def). FDB's retrieve() `param` +# key takes these numeric codes, not GRIB shortNames -- there's no lookup +# table for this shipped anywhere, so these were discovered empirically by +# decoding real archived messages under the correct ECCODES_DEFINITION_PATH +# (eccodes-cosmo-mars/definitions) and reading each field's `paramId`. To add +# a new one: fetch any archived field for that param, print +# `eccodes.codes_get(gid, "paramId")` with that definitions path active. +_FDB_PARAM_IDS: dict[str, str] = { + # levtype=sfc + "PS": "500000", + "PMSL": "500002", + "T_2M": "500011", + "TD_2M": "500017", + "U_10M": "500027", + "V_10M": "500029", + "TOT_PREC": "500041", + # levtype=pl + "FI": "500006", + "T": "500014", + "U": "500028", + "V": "500030", + "QV": "500035", +} +_FDB_PL_PARAMS = {"FI", "T", "U", "V", "QV"} + +# The schema's `timespan` key defaults to "none" when omitted from a request +# -- fine for instantaneous fields, but accumulated-from-start fields like +# TOT_PREC are archived with timespan="fs" and simply don't match (retrieve() +# silently returns nothing at all, for any step) if timespan isn't set +# explicitly. Confirmed empirically: TOT_PREC values this dict doesn't list +# default to "none". +_FDB_PARAM_TIMESPAN: dict[str, str] = { + "TOT_PREC": "fs", +} + +# _FDB_GLOBAL_MODEL ("varda-single-g", IFS-origin global grid) is archived +# under standard WMO/ECMWF param ids, not VARDA's custom local table -- see +# the `modifiers.patches` block in +# resources/inference/configs/sgm-temporal-downscaler-global_trimedge_multi-fdb.yaml, +# which explicitly pairs each ICON name with its archived (param, shortName). +# No pressure-level params mapped here: nothing in this codebase currently +# requests one from the global grid. +_FDB_GLOBAL_PARAM_IDS: dict[str, str] = { + "PS": "134", + "PMSL": "151", + "T_2M": "167", + "TD_2M": "168", + "U_10M": "165", + "V_10M": "166", + "TOT_PREC": "228", + "FR_LAND": "172", +} +_FDB_GLOBAL_PL_PARAMS: set[str] = set() +_FDB_GLOBAL_PARAM_TIMESPAN: dict[str, str] = { + "TOT_PREC": "fs", +} + +# Per-model (param_ids, pl_params, param_timespan) registries, keyed by the +# `model` value used in _fetch_fdb_fieldlist's requests. +_FDB_MODEL_REGISTRY: dict[str, tuple[dict[str, str], set[str], dict[str, str]]] = { + _FDB_MODEL: (_FDB_PARAM_IDS, _FDB_PL_PARAMS, _FDB_PARAM_TIMESPAN), + _FDB_GLOBAL_MODEL: (_FDB_GLOBAL_PARAM_IDS, _FDB_GLOBAL_PL_PARAMS, _FDB_GLOBAL_PARAM_TIMESPAN), +} + XARRAY_ENGINE_PROFILE = { "ensure_dims": ["z", "number", "step", "forecast_reference_time"], "global_attrs": [{"institution": "MeteoSwiss"}, {"Conventions": "CF-1.8"}], @@ -538,8 +601,18 @@ def fieldlist_to_xarray(fieldlist) -> xr.Dataset: ds = xr.Dataset() if len(fieldlist) == 0: return ds - for level_type_group in fieldlist.group_by("typeOfLevel"): - level_type = level_type_group[0].metadata("typeOfLevel") + # FieldList.group_by("typeOfLevel") is broken on the pinned earthkit-data + # (>=1.0,<1.1): it always yields a single group containing every field, + # regardless of how many distinct typeOfLevel values are actually present + # (reproduced in BUG_REPORT.md). Group by hand instead, using per-field + # .metadata() reads (reliable) rather than .group_by()/.sel() with the raw + # key (also unreliable on this version -- only the composite semantic keys + # like "parameter.variable" work via .sel()). + fields_by_level_type: dict[str, list] = {} + for fld in fieldlist: + fields_by_level_type.setdefault(fld.metadata("typeOfLevel"), []).append(fld) + for level_type, fields in fields_by_level_type.items(): + level_type_group = ekd.FieldList.from_fields(fields) profile = XARRAY_ENGINE_PROFILE | variable_name_profile(level_type) _ds = level_type_group.to_xarray(**profile, allow_holes=True) # add_valid_time_coord=True (earthkit's own convenience for this) isn't @@ -693,64 +766,106 @@ def fdb_has_data(fdb_root: str, reftime: datetime, model: str = _FDB_MODEL) -> b return bool(entries) -def _fetch_fdb_grib_bytes( - fdb_root: str, reftime: datetime, steps: list[int], model: str | None = None -) -> bytes: - """Retrieve raw GRIB bytes from an FDB store for one date/time and lead times. +def _fdb_config(fdb_root: str) -> dict: + return { + "type": "local", + "engine": "toc", + "schema": str(_FDB_SCHEMA.resolve()), + "spaces": [{"handler": "Default", "roots": [{"path": fdb_root}]}], + } + - `fdb.retrieve()` requires a single value per request for whatever keys - identify a distinct *block* (class/stream/type/expver/model/levtype) -- - unlike `fdb.list()`, which can wildcard across them -- so this lists first - to discover the blocks actually archived for this date/time/steps (e.g. one - per `output.tee` `fdb:` target), then issues one `retrieve()` per block and - concatenates their bytes (concatenated GRIB messages are themselves a valid - GRIB stream, readable as one file). +_fdb_identity_cache: dict[tuple, dict[str, tuple]] = {} - `model` restricts discovery to that exact value (VARDA-SINGLE archives - multiple grids under distinct model values, see `_FDB_MODEL`/ - `_FDB_GLOBAL_MODEL`); pass `None` to include every archived model instead. - Returns `b""` if nothing matches -- callers decide whether that's an error. + +def _discover_fdb_identity(fdb_root: str, model: str, date: str, time_: str) -> dict[str, tuple]: + """Return ``{levtype: (class, stream, type, expver)}`` actually archived + for this date/time/model -- these four keys vary by experiment/deployment + (e.g. class=od for a real operational archive vs ai/rd here), unlike + param/step, which callers always specify explicitly and so never need + discovering. `list()` is metadata-only (no data transfer) and fast, but + this is still cached per (fdb_root, model, date, time) since a single + `load_forecast_data` call looks it up twice (the caller's own params, then + the elevation fallback's FIS/FI lookup). """ import pyfdb + cache_key = (fdb_root, model, date, time_) + if cache_key in _fdb_identity_cache: + return _fdb_identity_cache[cache_key] + + with pyfdb.FDB(_fdb_config(fdb_root)) as fdb: + entries = list(fdb.list({"date": date, "time": time_, "model": model}, level=2)) + identities = {} + for entry in entries: + key = entry.combined_key() + identities[key["levtype"]] = (key["class"], key["stream"], key["type"], key["expver"]) + _fdb_identity_cache[cache_key] = identities + return identities + + +def _fetch_fdb_fieldlist( + fdb_root: str, reftime: datetime, steps: list[int], params: list[str], model: str = _FDB_MODEL, +): + """Fetch exactly the requested fields from FDB -- one `retrieve()` per + (levtype, timespan) group, each requesting exactly the caller's params via + FDB's own `param` request key (mapped through `_FDB_MODEL_REGISTRY`), + never "every param this block happens to have" discovered via `list()` + first. + + Uses earthkit-data's native `"fdb"` source (`stream=False`, which + retrieves via pyfdb and caches the result to a local temp file) rather + than calling pyfdb directly and hand-rolling our own tempfile/byte + concatenation -- earthkit already does exactly that, correctly (verified + directly against this archive; an earlier version of this function + claimed earthkit's `"fdb"` source was buggy, which doesn't reproduce on + the currently pinned earthkit-data). + + Returns `None` if nothing is archived for this date/time/model at all. + Raises `ValueError` for any param without a known FDB code rather than + silently fetching nothing for it, or for an unregistered `model` -- + see `_FDB_MODEL_REGISTRY`. + """ + if model not in _FDB_MODEL_REGISTRY: + raise ValueError(f"No FDB param registry for model={model!r} -- add one to _FDB_MODEL_REGISTRY.") + param_ids_table, pl_params, timespan_table = _FDB_MODEL_REGISTRY[model] + + unknown = [p for p in params if p not in param_ids_table] + if unknown: + raise ValueError( + f"No known FDB param id for {unknown} (model={model!r}) -- add it " + "to the relevant table in _FDB_MODEL_REGISTRY." + ) + date = reftime.strftime("%Y%m%d") time_ = reftime.strftime("%H%M") - config = { - "type": "local", - "engine": "toc", - "schema": str(_FDB_SCHEMA.resolve()), - "spaces": [{"handler": "Default", "roots": [{"path": fdb_root}]}], - } - # Keys identifying a distinct archived block; retrieve() needs a single value - # for each of these (list() can wildcard across them, retrieve() can't). - block_keys = ("class", "stream", "type", "expver", "model", "levtype") - list_request = {"date": date, "time": time_, "step": [str(s) for s in steps]} - if model is not None: - list_request["model"] = model + identities = _discover_fdb_identity(fdb_root, model, date, time_) + if not identities: + return None - with pyfdb.FDB(config) as fdb: - entries = list(fdb.list(list_request, level=3)) - if not entries: - return b"" - - blocks: dict[tuple, dict[str, set]] = {} - for entry in entries: - key = entry.combined_key() - block_key = tuple(key[bkey] for bkey in block_keys) - dims = blocks.setdefault(block_key, {}) - for k, v in key.items(): - if v not in (None, "") and k not in block_keys: - dims.setdefault(k, set()).add(v) - - buf = io.BytesIO() - for block_key, dims in blocks.items(): - request = dict(zip(block_keys, block_key)) - request.update({k: sorted(v) for k, v in dims.items()}) - handle = fdb.retrieve(request) - handle.open() - buf.write(handle.read()) - handle.close() - return buf.getvalue() + by_group: dict[tuple[str, str], list[str]] = {} + for p in params: + levtype = "pl" if p in pl_params else "sfc" + timespan = timespan_table.get(p, "none") + by_group.setdefault((levtype, timespan), []).append(param_ids_table[p]) + + config = _fdb_config(fdb_root) + fields = [] + for (levtype, timespan), param_ids in by_group.items(): + if levtype not in identities: + continue + cls_, stream, type_, expver = identities[levtype] + request = { + "class": cls_, "stream": stream, "type": type_, "expver": expver, + "model": model, "date": date, "time": time_, "levtype": levtype, + "step": [str(s) for s in steps], "param": sorted(set(param_ids)), + "timespan": timespan, + } + src = ekd.from_source("fdb", request, config=config, stream=False) + fl = src.to_fieldlist() if hasattr(src, "to_fieldlist") else src + fields.extend(fl) + + return ekd.FieldList.from_fields(fields) if fields else None def _load_forecast_data_from_fdb( @@ -761,34 +876,17 @@ def _load_forecast_data_from_fdb( External callers should use :func:`load_forecast_data` (root prefixed with ``"fdb:"``), which handles param expansion, step expansion, IC synthesis, and disaggregation, same as the GRIB-file path. Restricted to `_FDB_MODEL` - (the regional grid) for parity with `_collect_ml_grib_files`'s local-file - equivalent -- see the comment on `_FDB_MODEL`. - - The retrieved bytes are written to a temp GRIB file and reopened via the - existing local-file path (:func:`load_from_grib_file`) rather than decoded - in memory: earthkit-data's own `"fdb"` source has a bug (it never opens the - `pyfdb` `DataHandle` before reading -- raises `IOProblemError`), and - fieldlists built via `.sel()` / `FieldList.from_fields()` don't behave like - file-backed ones for metadata batch access. Writing to disk and reopening - as a genuine file sidesteps both and reuses the exact same, already-tested - GRIB-loading path used for local files. + (the regional grid, which uses VARDA's own ICON-convention param names + directly -- no IFS-alias expansion needed here, unlike the local-grib + path, which also has to handle the global forecaster's IFS-native output). """ - data = _fetch_fdb_grib_bytes(fdb_root, reftime, steps, model=_FDB_MODEL) - if not data: + fieldlist = _fetch_fdb_fieldlist(fdb_root, reftime, steps, params, model=_FDB_MODEL) + if fieldlist is None: raise FileNotFoundError( f"No FDB data found for date={reftime:%Y%m%d} time={reftime:%H%M} " f"model={_FDB_MODEL} in {fdb_root}" ) - - with tempfile.NamedTemporaryFile(suffix=".grib") as tmp: - tmp.write(data) - tmp.flush() - params_extended = list( - {p for p in params} | {_ICON_TO_IFS[p] for p in params if p in _ICON_TO_IFS} - ) - ds = load_from_grib_file(tmp.name, {"parameter.variable": params_extended}) - - return _rename_ifs_to_icon(ds) + return fieldlist_to_xarray(fieldlist) def _jretrieve_df_to_xarray(df, short_names, catalog) -> xr.Dataset: @@ -1515,7 +1613,11 @@ def load_forecast_data( # before falling back to ICON grid constants lookup. try: fis_ds = _load_forecast_data_from_fdb(fdb_root, reftime, [0], ["FIS"]) - except FileNotFoundError: + except (FileNotFoundError, ValueError): + # ValueError: "FIS" (surface geopotential) has no known FDB param + # id -- it's simply never archived for this model, confirmed by + # direct inspection; this probe is expected to come up empty here + # and fall through to the ICON-constants lookup below. fis_ds = None if fis_ds is not None and "FIS" in fis_ds: elevation = ekdv.geopotential_height_from_geopotential( diff --git a/src/plotting/compat.py b/src/plotting/compat.py index 07ff390a..60947f3c 100644 --- a/src/plotting/compat.py +++ b/src/plotting/compat.py @@ -6,7 +6,7 @@ import geopandas as gpd import numpy as np from shapely.geometry import MultiPoint -from data_input import load_from_grib_file, _fetch_fdb_grib_bytes, _FDB_MODEL, _FDB_GLOBAL_MODEL +from data_input import load_from_grib_file, _fetch_fdb_fieldlist, _FDB_MODEL, _FDB_GLOBAL_MODEL PARAMS_MAP = { @@ -116,7 +116,8 @@ def load_state_from_fdb( """FDB equivalent of :func:`load_state_from_grib` for one lead time. Fetches the regional (`_FDB_MODEL`) and global (`_FDB_GLOBAL_MODEL`) grids - from FDB and writes them to temp files using the same naming convention + from FDB -- exactly the requested params, via `_fetch_fdb_fieldlist` -- + and writes them to temp files using the same naming convention `load_state_from_grib` already expects locally (a primary file plus an `ifs-.grib` companion in the same directory), then delegates to it unchanged -- reusing its state-dict construction and global-grid @@ -125,18 +126,29 @@ def load_state_from_fdb( load_forecast_data` sidesteps by dropping the global grid entirely isn't a problem here). """ - regional = _fetch_fdb_grib_bytes(fdb_root, reftime, [step], model=_FDB_MODEL) - if not regional: + regional = _fetch_fdb_fieldlist(fdb_root, reftime, [step], paramlist, model=_FDB_MODEL) + if regional is None: raise FileNotFoundError( f"No FDB data found for step={step} model={_FDB_MODEL} in {fdb_root}" ) - global_ = _fetch_fdb_grib_bytes(fdb_root, reftime, [step], model=_FDB_GLOBAL_MODEL) + try: + global_ = _fetch_fdb_fieldlist(fdb_root, reftime, [step], paramlist, model=_FDB_GLOBAL_MODEL) + except ValueError: + # Some requested params may have no known id for the global grid + # (_FDB_GLOBAL_PARAM_IDS only covers what's actually been needed from + # it) -- tolerate that the same way as "nothing archived" below, + # rather than failing the whole frame over a companion grid. + global_ = None with tempfile.TemporaryDirectory() as tmp_dir: primary = Path(tmp_dir) / "frame.grib" - primary.write_bytes(regional) - if global_: - (Path(tmp_dir) / "ifs-frame.grib").write_bytes(global_) + with open(primary, "wb") as f: + for fld in regional: + f.write(fld.message()) + if global_ is not None: + with open(Path(tmp_dir) / "ifs-frame.grib", "wb") as f: + for fld in global_: + f.write(fld.message()) return load_state_from_grib(primary, paramlist=paramlist) From 528bbd9ef4d20d5890d598f70a966eedcd361b82 Mon Sep 17 00:00:00 2001 From: Carlos Osuna Date: Fri, 16 Oct 2026 14:59:37 +0200 Subject: [PATCH 07/21] perf: lazy-load heavy plotting deps out of plotting/__init__ cartopy, earthkit.plots and matplotlib.tri were imported eagerly at module load, forcing that whole stack onto any import of plotting.compat even for pure data-loading use (e.g. dashboard/meteogram reads). Move the viz-only content (DOMAINS, get_projection, StatePlotter) into plotting/viz.py and re-export it lazily via module-level __getattr__, so importing plotting only pulls in numpy until something viz-specific is actually touched. --- src/plotting/__init__.py | 295 +++------------------------------------ src/plotting/viz.py | 293 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 310 insertions(+), 278 deletions(-) create mode 100644 src/plotting/viz.py diff --git a/src/plotting/__init__.py b/src/plotting/__init__.py index 63b6be12..f835914d 100644 --- a/src/plotting/__init__.py +++ b/src/plotting/__init__.py @@ -1,286 +1,25 @@ -from contextlib import contextmanager -from pathlib import Path +"""Lightweight package root. Deliberately avoids importing anything heavy +(cartopy, earthkit.plots, matplotlib) at module level: Python always runs a +package's __init__.py before any of its submodules, so an eager import here +would force that whole visualization stack to load even for a submodule like +plotting.compat, which only loads data and never renders anything. + +DOMAINS/get_projection/StatePlotter (plotting.viz's actual content) are +re-exported lazily via module __getattr__ (PEP 562), so `from plotting import +DOMAINS` still works exactly as before for callers that do plot -- they just +pay the cartopy/earthkit.plots import cost at that point instead of always. +""" -import cartopy.crs as ccrs -import earthkit.plots as ekp import numpy as np -from matplotlib.tri import Triangulation State = dict[str, np.ndarray | dict[str, np.ndarray]] -_PROJECTIONS: dict[str, ccrs.Projection] = { - "platecarree": ccrs.PlateCarree(), - "orthographic": ccrs.Orthographic(central_longitude=5.0, central_latitude=45.0), - # added some pojections to test the behaviour, can be deleted later - "rotatedlatlon": ccrs.RotatedPole(pole_longitude=-170.0, pole_latitude=43.0), - "azimuthalequidist": ccrs.AzimuthalEquidistant(), -} +_LAZY_VIZ_NAMES = ("DOMAINS", "get_projection", "StatePlotter") -# Mapping of region names to their geographic extent and projection -# extent [lon_min, lon_max, lat_min, lat_max] in PlateCarree coordinates -def get_projection(name: str) -> "ccrs.Projection": - """Look up a projection by name.""" - if name not in _PROJECTIONS: - raise ValueError( - f"Unknown projection {name!r}. Available: {list(_PROJECTIONS)}" - ) - return _PROJECTIONS[name] +def __getattr__(name): + if name in _LAZY_VIZ_NAMES: + from . import viz - -DOMAINS = { - "globe": { - "extent": None, # full globe view - "projection": _PROJECTIONS["orthographic"], - }, - "europe": { - "extent": [-16.0, 25.0, 30.0, 65.0], - "projection": _PROJECTIONS["orthographic"], - }, - "centraleurope": { - "extent": [-1.5, 18, 41.5, 51], - "projection": _PROJECTIONS["orthographic"], - }, - "icon-ch": { - "extent": [0, 17.5, 40.5, 53.0], - "projection": _PROJECTIONS["orthographic"], - }, - "alps": { - "extent": [4.5, 16.0, 43.5, 48.5], - "projection": _PROJECTIONS["orthographic"], - }, - "radar": { - "extent": [3.2, 12.5, 43.6, 49.4], - "projection": _PROJECTIONS["orthographic"], - }, - "switzerland": { - "extent": [5.6, 10.8, 45.0, 48.6], - "projection": _PROJECTIONS["orthographic"], - }, -} - - -class StatePlotter: - """A class to plot state fields on various DOMAINS.""" - - def __init__( - self, - lon: np.ndarray, - lat: np.ndarray, - out_dir: Path, - ): - """Initialize the StatePlotter object. - - Latitudes and longitudes are passed during initialization so that - the triangulation can be computed once and reused for all plots. - The reference coordinate system of lon/lat is assumed to be PlateCarree() - currently. - - Parameters - ---------- - lon : np.ndarray - The longitudes of the grid. - lat : np.ndarray - The latitudes of the grid. - out_dir : Path - The output directory to save the plots. - """ - - self.lon = lon - self.lat = lat - out_dir.mkdir(exist_ok=True, parents=True) - self.out_dir = out_dir - - self.tri = Triangulation(self.lon, self.lat) - - def init_geoaxes( - self, - projection: ccrs.Projection, - bbox: list[float] | None, - nrows: int = 1, - ncols: int = 1, - size: tuple[float] | None = None, - name: str | None = None, - ) -> ekp.Figure: - """Initialize a figure with GeoAxes for plotting fields. - - Parameters - ---------- - projection : cartopy.crs.Projection - The projection used for the region. - bbox : list[float] or None - The bounding box [lon_min, lon_max, lat_min, lat_max] in PlateCarree coordinates. - If None, the full projection extent is used. - name : str, optional - The name of the region. - nrows : int, optional - The number of rows in the figure, by default 1. - ncols : int, optional - The number of columns in the figure, by default 1. - size : tuple - size of the figure in inches - - Returns - ------- - earthkit.plots.Figure - The figure object. - """ - if bbox is not None and len(bbox) != 4: - raise ValueError( - "bbox must be a list of four floats [lon_min, lon_max, lat_min, lat_max]" - ) - - ekp_fig = ekp.Figure( - crs=projection, - domain=bbox, - rows=nrows, - columns=ncols, - size=size, - ) - self.fig = ekp_fig - return ekp_fig - - def plot_field( - self, - subplot: ekp.Map, - field: np.ndarray, - style: ekp.styles.Style | None = None, - colorbar: bool = True, - title: str | None = None, - gridline_labels: bool = True, - **kwargs, - ): - """Plot a field on a Map object. - - Parameters - ---------- - subplot : earthkit.plots.Map - The Map subplot object to plot on. - field : np.ndarray - The field to plot. - style : ekp.styles.Style, optional - Earthkit.plots style for the map plot. - colorbar : bool - Whether to plot a colorbar, by default True. - title: str, optional - Map subplot title. - gridline_labels : bool - Whether to draw lat/lon degree labels on the gridlines, by - default True. Set to False for views whose center rotates - between frames (e.g. a rotating globe animation), where the - labels' varying width would otherwise make the map's position - shift from frame to frame. - kwargs : dict - Additional keyword arguments to pass to ax.tripcolor, including cmap, - vmin, vmax, etc. - """ - - proj = subplot._crs - # transform data coordinates to map coordinate reference system outside - # of the plotting function is a lot faster than letting tricontourf or - # tripcolor handle it in general, but not sure if using earthkit - # removed for now to simplify the workflow - if isinstance(proj, ccrs.Orthographic): - triang, mask = self._orthographic_tri(proj) - else: - triang, mask = self.tri, slice(None, None) - x, y = triang.x, triang.y - # TODO: this is hardcoded for when the initial state has two timesteps - # need to ditch this later - field = field[mask] - field = field[-1] if field.ndim == 2 else field.squeeze() - finite = np.isfinite(field) - # TODO: clip data to domain would make plotting faster (especially tripcolor) - # tried using Map.domain.extract() but too memory heavy (probably uses - # meshgrid in the background), implement clipping with e.g. - # points = np.column_stack((lon,lat)); mask = domain_polygon.contains_points(points) - # would work but needs handling domain crossing dateline etc. - - # TODO: tricontourf/tripcolor can handle a Triangulation when used directly, - # for some reason this doesn not work when using it with earthkit-plot, - # guess is that the earthkit-plots check throwing: - # "ValueError: x and y arrays must have the same length" is incorrect, - # therefore using x,y here - # subplot.tripcolor( # also works but is slower - # have to overwrite _plot_kwargs to avoid earthkit-plots trying to pass transform - # PlateCarree based on NumpySource - - # Normalize style and color-related kwargs - style_to_use, plot_kwargs = self._prepare_plot_kwargs(style, kwargs) - - # Temporarily suppress earthkit-plots internal source-based kwargs - with self._temporary_plot_kwargs_override(subplot): - subplot.tricontourf( - x=x[finite], - y=y[finite], - z=field[finite], - style=style_to_use, - transform=proj, - **plot_kwargs, - ) # for earthkit.plots to work properly cmap and norm are needed here - # TODO: gridlines etc would be nicer to have in the init, but I didn't get - # them to overlay the plot layer - - if gridline_labels: - subplot.standard_layers() - else: - subplot.land() - subplot.coastlines() - subplot.borders() - subplot.gridlines(draw_labels=False) - - if colorbar: - subplot.legend() - if title: - subplot.title(title) - - def _prepare_plot_kwargs( - self, - style: ekp.styles.Style | None, - kwargs: dict, - ) -> tuple[ekp.styles.Style | None, dict]: - """Return a cleaned style and plot kwargs without mutating the input.""" - plot_kwargs = dict(kwargs) - - # Discrete colors mode: if explicit 'colors' provided, drop cmap - colors = plot_kwargs.get("colors", None) - if colors is not None: - plot_kwargs.pop("cmap", None) - plot_kwargs.setdefault( - "no_style", True - ) # avoid interpolation being performed by earthkit-plots resulting in an error - return style, plot_kwargs - - # Continuous mode: remove None entries so configure_style doesn't call - # style.with_overrides(colors=None, vmin=None, ...) and wipe the style's - # pre-configured values. - for key in ("colors", "levels", "cmap", "norm", "vmin", "vmax"): - if plot_kwargs.get(key) is None: - plot_kwargs.pop(key, None) - - return style, plot_kwargs - - @contextmanager - def _temporary_plot_kwargs_override(self, subplot: ekp.Map): - """Temporarily override internal _plot_kwargs to avoid transform issues.""" - has_attr = hasattr(subplot, "_plot_kwargs") - old = getattr(subplot, "_plot_kwargs", None) - subplot._plot_kwargs = lambda source: {} - try: - yield - finally: - if has_attr: - subplot._plot_kwargs = old - else: - try: - delattr(subplot, "_plot_kwargs") - except Exception: - pass - - def _orthographic_tri( - self, proj: ccrs.Projection - ) -> tuple[Triangulation, np.ndarray]: - """Compute the triangulation for an orthographic-family projection.""" - x, y, _ = proj.transform_points(ccrs.PlateCarree(), self.lon, self.lat).T - mask = ~(np.isnan(x) | np.isnan(y)) - return Triangulation(x[mask], y[mask]), mask + return getattr(viz, name) + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/src/plotting/viz.py b/src/plotting/viz.py new file mode 100644 index 00000000..48ca4a44 --- /dev/null +++ b/src/plotting/viz.py @@ -0,0 +1,293 @@ +"""Visualization helpers (projections, domains, map rendering). + +Split out of plotting/__init__.py so that importing it (cartopy, +earthkit.plots, matplotlib) is only paid by callers that actually plot -- +plotting/__init__.py re-exports DOMAINS/get_projection/StatePlotter lazily +(see its module __getattr__) so that importing a sibling submodule like +plotting.compat (pure data loading, no rendering) never triggers this. +""" + +from contextlib import contextmanager +from pathlib import Path + +import cartopy.crs as ccrs +import earthkit.plots as ekp +import numpy as np +from matplotlib.tri import Triangulation + +_PROJECTIONS: dict[str, ccrs.Projection] = { + "platecarree": ccrs.PlateCarree(), + "orthographic": ccrs.Orthographic(central_longitude=5.0, central_latitude=45.0), + # added some pojections to test the behaviour, can be deleted later + "rotatedlatlon": ccrs.RotatedPole(pole_longitude=-170.0, pole_latitude=43.0), + "azimuthalequidist": ccrs.AzimuthalEquidistant(), +} + + +# Mapping of region names to their geographic extent and projection +# extent [lon_min, lon_max, lat_min, lat_max] in PlateCarree coordinates +def get_projection(name: str) -> "ccrs.Projection": + """Look up a projection by name.""" + if name not in _PROJECTIONS: + raise ValueError( + f"Unknown projection {name!r}. Available: {list(_PROJECTIONS)}" + ) + return _PROJECTIONS[name] + + +DOMAINS = { + "globe": { + "extent": None, # full globe view + "projection": _PROJECTIONS["orthographic"], + }, + "europe": { + "extent": [-16.0, 25.0, 30.0, 65.0], + "projection": _PROJECTIONS["orthographic"], + }, + "centraleurope": { + "extent": [-1.5, 18, 41.5, 51], + "projection": _PROJECTIONS["orthographic"], + }, + "icon-ch": { + "extent": [0, 17.5, 40.5, 53.0], + "projection": _PROJECTIONS["orthographic"], + }, + "alps": { + "extent": [4.5, 16.0, 43.5, 48.5], + "projection": _PROJECTIONS["orthographic"], + }, + "radar": { + "extent": [3.2, 12.5, 43.6, 49.4], + "projection": _PROJECTIONS["orthographic"], + }, + "switzerland": { + "extent": [5.6, 10.8, 45.0, 48.6], + "projection": _PROJECTIONS["orthographic"], + }, +} + + +class StatePlotter: + """A class to plot state fields on various DOMAINS.""" + + def __init__( + self, + lon: np.ndarray, + lat: np.ndarray, + out_dir: Path, + ): + """Initialize the StatePlotter object. + + Latitudes and longitudes are passed during initialization so that + the triangulation can be computed once and reused for all plots. + The reference coordinate system of lon/lat is assumed to be PlateCarree() + currently. + + Parameters + ---------- + lon : np.ndarray + The longitudes of the grid. + lat : np.ndarray + The latitudes of the grid. + out_dir : Path + The output directory to save the plots. + """ + + self.lon = lon + self.lat = lat + out_dir.mkdir(exist_ok=True, parents=True) + self.out_dir = out_dir + + self.tri = Triangulation(self.lon, self.lat) + + def init_geoaxes( + self, + projection: ccrs.Projection, + bbox: list[float] | None, + nrows: int = 1, + ncols: int = 1, + size: tuple[float] | None = None, + name: str | None = None, + ) -> ekp.Figure: + """Initialize a figure with GeoAxes for plotting fields. + + Parameters + ---------- + projection : cartopy.crs.Projection + The projection used for the region. + bbox : list[float] or None + The bounding box [lon_min, lon_max, lat_min, lat_max] in PlateCarree coordinates. + If None, the full projection extent is used. + name : str, optional + The name of the region. + nrows : int, optional + The number of rows in the figure, by default 1. + ncols : int, optional + The number of columns in the figure, by default 1. + size : tuple + size of the figure in inches + + Returns + ------- + earthkit.plots.Figure + The figure object. + """ + if bbox is not None and len(bbox) != 4: + raise ValueError( + "bbox must be a list of four floats [lon_min, lon_max, lat_min, lat_max]" + ) + + ekp_fig = ekp.Figure( + crs=projection, + domain=bbox, + rows=nrows, + columns=ncols, + size=size, + ) + self.fig = ekp_fig + return ekp_fig + + def plot_field( + self, + subplot: ekp.Map, + field: np.ndarray, + style: ekp.styles.Style | None = None, + colorbar: bool = True, + title: str | None = None, + gridline_labels: bool = True, + **kwargs, + ): + """Plot a field on a Map object. + + Parameters + ---------- + subplot : earthkit.plots.Map + The Map subplot object to plot on. + field : np.ndarray + The field to plot. + style : ekp.styles.Style, optional + Earthkit.plots style for the map plot. + colorbar : bool + Whether to plot a colorbar, by default True. + title: str, optional + Map subplot title. + gridline_labels : bool + Whether to draw lat/lon degree labels on the gridlines, by + default True. Set to False for views whose center rotates + between frames (e.g. a rotating globe animation), where the + labels' varying width would otherwise make the map's position + shift from frame to frame. + kwargs : dict + Additional keyword arguments to pass to ax.tripcolor, including cmap, + vmin, vmax, etc. + """ + + proj = subplot._crs + # transform data coordinates to map coordinate reference system outside + # of the plotting function is a lot faster than letting tricontourf or + # tripcolor handle it in general, but not sure if using earthkit + # removed for now to simplify the workflow + if isinstance(proj, ccrs.Orthographic): + triang, mask = self._orthographic_tri(proj) + else: + triang, mask = self.tri, slice(None, None) + x, y = triang.x, triang.y + # TODO: this is hardcoded for when the initial state has two timesteps + # need to ditch this later + field = field[mask] + field = field[-1] if field.ndim == 2 else field.squeeze() + finite = np.isfinite(field) + # TODO: clip data to domain would make plotting faster (especially tripcolor) + # tried using Map.domain.extract() but too memory heavy (probably uses + # meshgrid in the background), implement clipping with e.g. + # points = np.column_stack((lon,lat)); mask = domain_polygon.contains_points(points) + # would work but needs handling domain crossing dateline etc. + + # TODO: tricontourf/tripcolor can handle a Triangulation when used directly, + # for some reason this doesn not work when using it with earthkit-plot, + # guess is that the earthkit-plots check throwing: + # "ValueError: x and y arrays must have the same length" is incorrect, + # therefore using x,y here + # subplot.tripcolor( # also works but is slower + # have to overwrite _plot_kwargs to avoid earthkit-plots trying to pass transform + # PlateCarree based on NumpySource + + # Normalize style and color-related kwargs + style_to_use, plot_kwargs = self._prepare_plot_kwargs(style, kwargs) + + # Temporarily suppress earthkit-plots internal source-based kwargs + with self._temporary_plot_kwargs_override(subplot): + subplot.tricontourf( + x=x[finite], + y=y[finite], + z=field[finite], + style=style_to_use, + transform=proj, + **plot_kwargs, + ) # for earthkit.plots to work properly cmap and norm are needed here + # TODO: gridlines etc would be nicer to have in the init, but I didn't get + # them to overlay the plot layer + + if gridline_labels: + subplot.standard_layers() + else: + subplot.land() + subplot.coastlines() + subplot.borders() + subplot.gridlines(draw_labels=False) + + if colorbar: + subplot.legend() + if title: + subplot.title(title) + + def _prepare_plot_kwargs( + self, + style: ekp.styles.Style | None, + kwargs: dict, + ) -> tuple[ekp.styles.Style | None, dict]: + """Return a cleaned style and plot kwargs without mutating the input.""" + plot_kwargs = dict(kwargs) + + # Discrete colors mode: if explicit 'colors' provided, drop cmap + colors = plot_kwargs.get("colors", None) + if colors is not None: + plot_kwargs.pop("cmap", None) + plot_kwargs.setdefault( + "no_style", True + ) # avoid interpolation being performed by earthkit-plots resulting in an error + return style, plot_kwargs + + # Continuous mode: remove None entries so configure_style doesn't call + # style.with_overrides(colors=None, vmin=None, ...) and wipe the style's + # pre-configured values. + for key in ("colors", "levels", "cmap", "norm", "vmin", "vmax"): + if plot_kwargs.get(key) is None: + plot_kwargs.pop(key, None) + + return style, plot_kwargs + + @contextmanager + def _temporary_plot_kwargs_override(self, subplot: ekp.Map): + """Temporarily override internal _plot_kwargs to avoid transform issues.""" + has_attr = hasattr(subplot, "_plot_kwargs") + old = getattr(subplot, "_plot_kwargs", None) + subplot._plot_kwargs = lambda source: {} + try: + yield + finally: + if has_attr: + subplot._plot_kwargs = old + else: + try: + delattr(subplot, "_plot_kwargs") + except Exception: + pass + + def _orthographic_tri( + self, proj: ccrs.Projection + ) -> tuple[Triangulation, np.ndarray]: + """Compute the triangulation for an orthographic-family projection.""" + x, y, _ = proj.transform_points(ccrs.PlateCarree(), self.lon, self.lat).T + mask = ~(np.isnan(x) | np.isnan(y)) + return Triangulation(x[mask], y[mask]), mask From 45900ab5040359bbd8f9285a6911f38fce0da36f Mon Sep 17 00:00:00 2001 From: Carlos Osuna Date: Fri, 16 Oct 2026 14:59:59 +0200 Subject: [PATCH 08/21] fix: squashfs-mount -s flag syntax and pin eccodes-cosmo-mars commit squashfs-mount v10.1.0 replaced the old positional "path:mountpoint ..." syntax with a single -s "path:mnt,path2:mnt2" flag; update both call sites (inference_check_fdb and inference_execute) accordingly, and make the .venv-mount cleanup trap tolerate it already being gone. Also pin the eccodes-cosmo-mars clone to commit d04363540b (pre "add a grid Name") instead of a bare --depth 1 branch clone: a later commit on varda-ext adds a mars.gridName concept that the current realtime-varda FDB schema doesn't account for, breaking archival with a "Keywords not used" error. Un-pin once the schema is updated to match. --- workflow/rules/inference.smk | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/workflow/rules/inference.smk b/workflow/rules/inference.smk index 91c047ee..1044abb9 100644 --- a/workflow/rules/inference.smk +++ b/workflow/rules/inference.smk @@ -123,10 +123,16 @@ rule inference_prepare_env: echo "[$(date)] Cloning eccodes-cosmo-mars (varda-ext branch) for MARS definitions..." mkdir -p "$VENV/data/share" - git clone --branch varda-ext --depth 1 \ + # Pinned to d043635 (pre "add a grid Name") rather than a bare + # --depth 1 branch clone: f9b4603 adds a mars.gridName concept + # that the current realtime-varda.schema doesn't account for, + # breaking FDB archival with a "Keywords not used" error. + # Un-pin once the schema is updated to match. + git clone --branch varda-ext \ git@github.com:meteoswiss/eccodes-cosmo-mars.git \ "$VENV/data/share/eccodes-cosmo-mars" - echo "[$(date)] Cloned eccodes-cosmo-mars definitions" + git -C "$VENV/data/share/eccodes-cosmo-mars" checkout d04363540bb24b1fd432dce2b1e65ba0f27ba591 + echo "[$(date)] Cloned eccodes-cosmo-mars definitions (pinned to d043635)" # metkitlib (pulled in by pyfdb) ships its own language.yaml, used to # validate fdb.retrieve() requests. Our custom `model` values (added @@ -293,9 +299,9 @@ checkpoint inference_check_fdb: # is a pre-provisioned system directory, but there's no such slot for a second # image, so the venv gets a job-local directory we create ourselves. VENV_MOUNT=$(mktemp -d) - trap "rmdir '$VENV_MOUNT'" EXIT + trap "rmdir '$VENV_MOUNT' 2>/dev/null || true" EXIT FDB_SQUASHFS=$(uenv image inspect --format='{{sqfs}}' {params.fdb_uenv}) - squashfs-mount {params.image_path}:$VENV_MOUNT "$FDB_SQUASHFS:/user-environment" -- bash -c ' + squashfs-mount -s "{params.image_path}:$VENV_MOUNT,$FDB_SQUASHFS:/user-environment" -- bash -c ' source '"$VENV_MOUNT"'/bin/activate export LD_LIBRARY_PATH=/user-environment/env/{params.fdb_view}/lib64:/user-environment/env/{params.fdb_view}/lib${{LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}} python {input.script} '"$GLOBAL_ARG"' \ @@ -568,7 +574,7 @@ FDBEOF # mount around srun is fine; two images at once are not. VENV_MOUNT="{params.workdir}/.venv-mount" mkdir -p "$VENV_MOUNT" - trap "rmdir '$VENV_MOUNT'" EXIT + trap "rmdir '$VENV_MOUNT' 2>/dev/null || true" EXIT MOUNT_ARGS=("{params.env_path}:$VENV_MOUNT") INNER_CMD="_run_inference $VENV_MOUNT" @@ -586,7 +592,7 @@ FDBEOF --time={resources.runtime} \ --gres={resources.gres} \ --ntasks={resources.ntasks} \ - squashfs-mount "${{MOUNT_ARGS[@]}}" -- bash -c "$INNER_CMD" + squashfs-mount -s "$(IFS=,; echo "${{MOUNT_ARGS[*]}}")" -- bash -c "$INNER_CMD" ) >{log} 2>&1 && touch {output.okfile} """ # fmt: on From 280f01007dd8978e3b80783d9d3b69e62a9d5653 Mon Sep 17 00:00:00 2001 From: Carlos Osuna Date: Fri, 16 Oct 2026 15:00:14 +0200 Subject: [PATCH 09/21] fix: tolerate anemoi-inference's benign native crash during teardown anemoi-inference, when using the copy-prognostic-from-forecaster post-processor (anemoi-plugins-meteoswiss), reliably crashes natively during process teardown strictly AFTER all forecast windows have been computed and written out -- glibc heap-corruption diagnostics (free(): invalid pointer, munmap_chunk(): invalid pointer, double free or corruption), with the exact message/signal (SIGABRT/134, SIGBUS/135) depending on memory layout at the time. Likely an eccodes/earthkit-data handle cleanup ordering bug tied to that filter's cached forecaster GRIB index being garbage-collected at interpreter shutdown. Verified empirically (including "Flushing FDB" x2 immediately preceding the crash, for FDB-backed runs) that all expected output is present and bit-correct despite the abort. Without this, the rule discarded valid output and forced a full re-run of expensive GPU inference for nothing. Detect the signature (no Python traceback, all forecast windows reached) and treat it as success instead. --- workflow/rules/inference.smk | 35 ++++++++++++++++++++++++++++++++++- 1 file changed, 34 insertions(+), 1 deletion(-) diff --git a/workflow/rules/inference.smk b/workflow/rules/inference.smk index 1044abb9..81adb584 100644 --- a/workflow/rules/inference.smk +++ b/workflow/rules/inference.smk @@ -557,7 +557,40 @@ FDBEOF CMD_ARGS+=(runner.parallel.cluster=slurm) fi - anemoi-inference run config.yaml "${{CMD_ARGS[@]}}" + set +e + anemoi-inference run config.yaml "${{CMD_ARGS[@]}}" 2>&1 | tee inference_stdout.log + local rc=${{PIPESTATUS[0]}} + set -e + if [ "$rc" -ne 0 ]; then + # anemoi-inference, when using the copy-prognostic-from-forecaster + # post-processor (anemoi-plugins-meteoswiss), is known to crash + # during native library teardown strictly AFTER all forecast windows + # have been computed and written out (free(): invalid pointer, + # munmap_chunk(): invalid pointer, or similar glibc heap-corruption + # diagnostics observed across runs -- likely an eccodes/earthkit-data + # handle cleanup ordering bug tied to + # that filter's cached forecaster GRIB index being garbage-collected + # at interpreter shutdown). The exact signal raised depends on memory + # layout at the time: observed as both SIGABRT (134) and SIGBUS (135) + # across otherwise-identical runs. Verified empirically (including + # "Flushing FDB" x2 immediately preceding the crash, for FDB-backed + # runs): all expected output steps are present and bit-correct on + # disk/FDB despite the abort. Treat this exact signature as success + # instead of discarding valid output and re-running (expensive) GPU + # inference for nothing. + local last_window + last_window=$(grep -oE 'window [0-9]+/[0-9]+' inference_stdout.log | tail -1 | awk '{{print $2}}') + local win_num=${{last_window%%/*}} + local win_total=${{last_window##*/}} + if [ "$rc" = "134" -o "$rc" = "135" ] \ + && grep -q -e "free(): invalid pointer" -e "munmap_chunk(): invalid pointer" -e "double free or corruption" inference_stdout.log \ + && ! grep -q "Traceback (most recent call last)" inference_stdout.log \ + && [ -n "$win_num" ] && [ "$win_num" = "$win_total" ]; then + echo "WARNING: anemoi-inference aborted during process teardown (exit $rc) after completing all $win_total forecast windows -- known benign native cleanup crash (see copy-prognostic-from-forecaster + eccodes handle GC), treating as success." >&2 + else + return "$rc" + fi + fi }} export -f _run_inference From ec6a1a72b691c66749da7e5a6363331b57bdb3e3 Mon Sep 17 00:00:00 2001 From: Carlos Osuna Date: Fri, 16 Oct 2026 15:00:25 +0200 Subject: [PATCH 10/21] fix: decompose aggregated param suffixes (TOT_PREC1/6) in frame plotting plot_forecast_frame.py received showcase-facing param names like "TOT_PREC1"/"TOT_PREC6" (base param + accumulation window encoded in a numeric suffix) and passed them straight through to GRIB loading and de-accumulation, which only ever matched the literal base name "TOT_PREC" -- every precipitation frame failed with "None of the requested params ['TOT_PREC1'] found". Decompose via the existing parse_aggregated_param helper, using the base name for loading/de-accumulation/colormap lookup and the original suffixed name only for filenames/titles. Same root cause in resolve_leadtimes() (common.smk): the step-0 skip for accumulated params (PARAMS_WITHOUT_STEP_ZERO_VALUE) only recognized the literal "TOT_PREC", not its suffixed showcase variants, so frames for leadtime 0 failed with "accumulated fields ... undefined at step 0" where they should have simply been skipped. --- workflow/rules/common.smk | 7 +++- workflow/scripts/plot_forecast_frame.py | 45 ++++++++++++++++++++----- 2 files changed, 43 insertions(+), 9 deletions(-) diff --git a/workflow/rules/common.smk b/workflow/rules/common.smk index 5342d795..29631f77 100644 --- a/workflow/rules/common.smk +++ b/workflow/rules/common.smk @@ -1,5 +1,6 @@ import logging import copy +import re from datetime import datetime, timedelta import yaml import hashlib @@ -518,6 +519,10 @@ def resolve_leadtimes(steps_spec, requested="all", param=None): ) valid = wanted & supported - if param in PARAMS_WITHOUT_STEP_ZERO_VALUE: + # Showcase-facing names encode an accumulation window as a numeric suffix + # (e.g. "TOT_PREC1", "TOT_PREC6"); strip it before checking membership so + # those aggregated variants are recognized the same as their base param. + base_param = re.sub(r"\d+$", "", param) if param else param + if param in PARAMS_WITHOUT_STEP_ZERO_VALUE or base_param in PARAMS_WITHOUT_STEP_ZERO_VALUE: valid = {lt for lt in valid if lt >= step} return sorted(valid) diff --git a/workflow/scripts/plot_forecast_frame.py b/workflow/scripts/plot_forecast_frame.py index 472611f5..cadc8237 100644 --- a/workflow/scripts/plot_forecast_frame.py +++ b/workflow/scripts/plot_forecast_frame.py @@ -11,6 +11,7 @@ from matplotlib.colors import Colormap import numpy as np +from data_input import parse_aggregated_param from plotting import DOMAINS from plotting import get_projection from plotting import StatePlotter @@ -124,10 +125,15 @@ def main(): init_time = args.date lead_time = args.leadtime param = args.param + # Showcase-facing param names encode the accumulation window in the suffix + # (e.g. "TOT_PREC1" -> base "TOT_PREC", 1h; "TOT_PREC6" -> 6h). Data loading, + # de-accumulation and colormap lookup all key off the base GRIB/FDB param + # name; `param` itself is kept only for filenames/titles. + base_param, agg_hours = parse_aggregated_param(param) regions = json.loads(args.regions_json) outdir = Path(args.outdir) outdir.mkdir(parents=True, exist_ok=True) - accu = args.accu + accu = agg_hours if agg_hours is not None else args.accu LOG.info( "Plotting forecast frame: param=%s, init_time=%s, lead_time=%s, regions=%s", @@ -137,12 +143,12 @@ def main(): list(regions.keys()), ) - if param == "SP_10M": + if base_param == "SP_10M": paramlist = ["U_10M", "V_10M"] - elif param == "SP": + elif base_param == "SP": paramlist = ["U", "V"] else: - paramlist = [param] + paramlist = [base_param] def _empty_frames(): for region_name in regions: @@ -165,6 +171,17 @@ def _empty_frames(): ) _empty_frames() return + except ValueError as exc: + # Accumulated params (e.g. TOT_PREC1/6) are undefined at lead times + # shorter than their accumulation window -- normally excluded + # upstream (resolve_leadtimes), but handled here too as a + # defensive fallback rather than crashing the whole rule. + LOG.warning( + "%s Creating empty placeholder frames.", + exc, + ) + _empty_frames() + return else: grib_dir = Path(args.input) # TODO: fix file pattern & globbing @@ -182,10 +199,22 @@ def _empty_frames(): return grib_file = Path(grib_file[0]) LOG.info("Loading grib file %s", grib_file) - state = load_state_from_grib(grib_file, paramlist=paramlist) + try: + state = load_state_from_grib(grib_file, paramlist=paramlist) + except ValueError as exc: + # Accumulated params (e.g. TOT_PREC1/6) are undefined at lead times + # shorter than their accumulation window -- normally excluded + # upstream (resolve_leadtimes), but handled here too as a + # defensive fallback rather than crashing the whole rule. + LOG.warning( + "%s Creating empty placeholder frames.", + exc, + ) + _empty_frames() + return # tp is accumulated from start of forecast; de-accumulate to get period [lt-accu, lt] - if param == "TOT_PREC": + if base_param == "TOT_PREC": prev_lt = lead_time - accu if prev_lt > 0: if is_fdb: @@ -212,7 +241,7 @@ def _empty_frames(): ) # Preprocess field once — shared across all region plots - field, units_override = preprocess_field(param, state) + field, units_override = preprocess_field(base_param, state) validtime = state["valid_time"].strftime("%Y%m%d%H%M") for region_name, region_cfg in regions.items(): @@ -250,7 +279,7 @@ def _empty_frames(): field, title=f"{param}, time: {validtime}", gridline_labels=not region_cfg.get("rotate", False), - **get_style(param, units_override, accu=accu), + **get_style(base_param, units_override, accu=accu), ) if len(state["lam_envelope"]) > 0: subplot.ax.add_geometries( From c95864e772bac60d76e6f48a68cd65ded6ca793e Mon Sep 17 00:00:00 2001 From: Carlos Osuna Date: Fri, 16 Oct 2026 15:00:40 +0200 Subject: [PATCH 11/21] fix: bump plot_forecast_frame Slurm runtime budget to 45m Under heavy concurrent load on the shared postproc partition (many plot_forecast_frame jobs running at once), a single region's render has been observed taking ~16m vs ~1.5m uncontended. The 20m budget was killing otherwise-healthy jobs via TIMEOUT -- a resource/contention issue, not a correctness bug -- discarding completed region renders along with it. --- workflow/rules/plot.smk | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/workflow/rules/plot.smk b/workflow/rules/plot.smk index 10ef02ec..d5ccecae 100644 --- a/workflow/rules/plot.smk +++ b/workflow/rules/plot.smk @@ -143,7 +143,11 @@ rule plot_forecast_frame: resources: slurm_partition="postproc", cpus_per_task=4, - runtime="20m", + # Bumped from 20m: under heavy concurrent load on the shared postproc + # partition (many plot_forecast_frame jobs at once), a single region's + # render has been observed taking ~16m vs ~1.5m uncontended -- 20m was + # killing otherwise-healthy jobs via TIMEOUT, not a correctness bug. + runtime="45m", params: grib_out_dir=lambda wc: str( (Path(OUT_ROOT) / f"data/runs/{wc.run_id}/{wc.init_time}/grib").resolve() From 1bf94b390489e0613b493001085b507b7d1d44c7 Mon Sep 17 00:00:00 2001 From: Carlos Osuna Date: Tue, 20 Oct 2026 20:45:01 +0200 Subject: [PATCH 12/21] adding policy for imagemagick --- resources/imagemagick/policy.xml | 21 +++++++++++++++++++++ workflow/rules/plot.smk | 8 ++++++++ 2 files changed, 29 insertions(+) create mode 100644 resources/imagemagick/policy.xml diff --git a/resources/imagemagick/policy.xml b/resources/imagemagick/policy.xml new file mode 100644 index 00000000..e75da125 --- /dev/null +++ b/resources/imagemagick/policy.xml @@ -0,0 +1,21 @@ + + + + + + + + + + + diff --git a/workflow/rules/plot.smk b/workflow/rules/plot.smk index d5ccecae..a2b3cf0a 100644 --- a/workflow/rules/plot.smk +++ b/workflow/rules/plot.smk @@ -227,6 +227,14 @@ rule make_forecast_animation: shell: """ FRAMES=$(for f in {input}; do [ -s "$f" ] && echo "$f"; done | tr '\\n' ' ') + # The system ImageMagick policy.xml caps "area" at 16KP (16 thousand + # pixels -- smaller than almost any real image) as a hardening + # measure; that's a hard ceiling the -limit CLI flag cannot raise, so + # assembling ~120 full-resolution frames into a GIF silently failed + # ("cache resources exhausted" / "TooManyExceptions"). Point + # MAGICK_CONFIGURE_PATH at our own policy.xml with generous (not + # unbounded) limits instead. + export MAGICK_CONFIGURE_PATH=$(realpath resources/imagemagick) convert -delay {params.delay} -loop 0 $FRAMES {output} """ From 1f535f163df1f4c3e8a8c2085f95d68f4f883cb3 Mon Sep 17 00:00:00 2001 From: Carlos Osuna Date: Wed, 21 Oct 2026 09:10:43 +0200 Subject: [PATCH 13/21] useand output key in the config to selectgrib or fdb --- config/varda-single-1.0-inference-only.yaml | 1 + src/data_input/__init__.py | 33 --------------- src/evalml/config.py | 26 +++++++++++- workflow/rules/common.smk | 6 +++ workflow/rules/inference.smk | 42 +++++++++++++++---- workflow/rules/plot.smk | 36 +++++----------- workflow/rules/verification.smk | 46 ++++++++------------- workflow/scripts/verification_scoremaps.py | 41 +++++++++--------- workflow/tools/config.schema.json | 20 +++++++++ 9 files changed, 136 insertions(+), 115 deletions(-) diff --git a/config/varda-single-1.0-inference-only.yaml b/config/varda-single-1.0-inference-only.yaml index 583faca3..91959060 100644 --- a/config/varda-single-1.0-inference-only.yaml +++ b/config/varda-single-1.0-inference-only.yaml @@ -15,6 +15,7 @@ runs: label: Varda-single-1.0 steps: 0/120/1 config: resources/inference/configs/sgm-temporal-downscaler-global_trimedge_multi-fdb.yaml + output: fdb extra_requirements: - anemoi-datasets==0.5.35 - /scratch/mch/cosuna/evalml_fdb/anemoi-inference diff --git a/src/data_input/__init__.py b/src/data_input/__init__.py index ee531e31..7045f681 100644 --- a/src/data_input/__init__.py +++ b/src/data_input/__init__.py @@ -733,39 +733,6 @@ def _rename_ifs_to_icon(ds: xr.Dataset) -> xr.Dataset: return ds.rename(ifs_rename) if ifs_rename else ds -def fdb_has_data(fdb_root: str, reftime: datetime, model: str = _FDB_MODEL) -> bool: - """Lightweight existence check: does `fdb_root` have any entry for this - date/time under `model`? - - Doesn't check completeness (see workflow/scripts/inference_check_fdb.py's - block/step/param accounting for that) -- for callers where completeness was - already established upstream (e.g. by a Snakemake dependency on this - reftime's inference_check_fdb/inference_execute okfile before this runs), - a bare presence check is enough to decide whether to route a reftime to - the FDB path at all. - """ - import pyfdb - - config = { - "type": "local", - "engine": "toc", - "schema": str(_FDB_SCHEMA.resolve()), - "spaces": [{"handler": "Default", "roots": [{"path": fdb_root}]}], - } - with pyfdb.FDB(config) as fdb: - entries = list( - fdb.list( - { - "date": reftime.strftime("%Y%m%d"), - "time": reftime.strftime("%H%M"), - "model": model, - }, - level=1, - ) - ) - return bool(entries) - - def _fdb_config(fdb_root: str) -> dict: return { "type": "local", diff --git a/src/evalml/config.py b/src/evalml/config.py index 9730d7d3..69ebf18c 100644 --- a/src/evalml/config.py +++ b/src/evalml/config.py @@ -1,5 +1,5 @@ from pathlib import Path -from typing import Dict, List, Any, ClassVar, FrozenSet, Optional, Union +from typing import Dict, List, Any, ClassVar, FrozenSet, Literal, Optional, Union from pydantic import BaseModel, Field, RootModel, field_validator, model_validator @@ -129,6 +129,15 @@ class RunConfig(BaseModel): False, description="If true, the ECCODES_DEFINITION_PATH environment variable will not be set to the COSMO local definitions.", ) + output: Literal["grib", "fdb"] = Field( + "grib", + description="Where this run's forecast output is stored, and therefore where " + "downstream reads (verification, meteograms, frames, scoremaps) should look for " + "it. Must match what the referenced inference `config`'s own output.tee actually " + "writes to (an 'fdb' tee block requires output: fdb here, and vice versa) -- this " + "is the single declared source of truth for backend routing; nothing probes for " + "data presence or parses the inference config to guess it.", + ) config: Dict[str, Any] | str model_config = {"extra": "forbid"} @@ -187,6 +196,21 @@ class TemporalDownscalerConfig(RunConfig): description="Configuration for the forecaster run that this temporal downscaler is based on.", ) + @model_validator(mode="after") + def _forecaster_output_must_be_grib(self) -> "TemporalDownscalerConfig": + # inference_prepare.py's prepare_temporal_downscaler() hardcodes a symlink to + # the forecaster's local grib directory to chain its output into the downscaler + # as input -- an FDB-backed forecaster would make that symlink point at an empty + # directory, failing confusingly deep inside anemoi-inference rather than here. + if self.forecaster is not None and self.forecaster.output == "fdb": + raise ValueError( + "temporal_downscaler.forecaster.output: fdb' is not supported -- the " + "downscaler reads its forecaster's prognostic state via a local grib " + "symlink (inference_prepare.py), so the forecaster must use output: grib " + "(the default). Only the temporal_downscaler's own output may be fdb." + ) + return self + class BaselineConfig(BaseModel): """Configuration for a single baseline to include in the verification.""" diff --git a/workflow/rules/common.smk b/workflow/rules/common.smk index 29631f77..5003665c 100644 --- a/workflow/rules/common.smk +++ b/workflow/rules/common.smk @@ -361,11 +361,17 @@ def run_specific_hash(run_config: dict, model_type: str) -> str: Changes to these fields create a new run_id (new outputs) but reuse the environment: - steps (lead times) - config YAML file contents (inference parameters) + - output (grib/fdb), only when non-default -- see below - For temporal downscalers: the forecaster's run_id (which run's outputs to read) """ configs_to_hash = [{"steps": run_config["steps"]}] with open(run_config["config"], "r") as f: configs_to_hash.append(yaml.safe_load(f)) + # Hashed only when set to the non-default "fdb": hashing it unconditionally would + # change run_id (and thus the output directory) for every existing run in the repo + # the moment this field shipped, making every already-computed grib run look stale. + if run_config.get("output", "grib") != "grib": + configs_to_hash.append({"output": run_config["output"]}) if model_type == "temporal_downscaler" and run_config.get("forecaster"): # run output depends on which forecaster RUN was used (not just the env) configs_to_hash.append(run_config["forecaster"].get("run_id")) diff --git a/workflow/rules/inference.smk b/workflow/rules/inference.smk index 81adb584..28b56081 100644 --- a/workflow/rules/inference.smk +++ b/workflow/rules/inference.smk @@ -255,8 +255,9 @@ checkpoint inference_check_fdb: time to be present in one root -- see inference_check_fdb.py for why no further MARS-key disambiguation (class/expver/model) is needed. Writes 'hit\\t' or 'miss\\t' to the status file. Runs inside the squashfs venv so that pyfdb and the - native FDB C library are available. Skipped entirely (writes 'miss') when this - run's environment has no FDB uenv bundled (profile.fdb.srun_uenv unset). + native FDB C library are available. Skipped entirely (writes 'miss') unless this + run's own config declares `output: fdb` (see the fail-fast validation below this + rule, which rejects output: fdb runs whose environment has no FDB uenv bundled). """ input: image=lambda wc: OUT_ROOT @@ -269,9 +270,7 @@ checkpoint inference_check_fdb: localrule: True params: image_path=lambda wc, input: str(Path(input.image).resolve()), - fdb_configured=lambda wc: bool( - _get_fdb_uenv_for_env(RUN_CONFIGS[wc.run_id]["env_id"])[0] - ), + fdb_configured=lambda wc: RUN_CONFIGS[wc.run_id]["output"] == "fdb", fdb_uenv=lambda wc: _get_fdb_uenv_for_env(RUN_CONFIGS[wc.run_id]["env_id"])[0], fdb_root_global=lambda wc: _get_fdb_roots(wc.run_id)[0], fdb_root_local=lambda wc: _get_fdb_roots(wc.run_id)[1], @@ -489,9 +488,7 @@ rule inference_execute: srun_prefix=lambda wc: get_resource(wc, "srun_prefix", "") or "", fdb_view=lambda wc: get_resource(wc, "srun_view", "realtime") or "realtime", fdb_uenv=lambda wc: _get_fdb_uenv_for_env(RUN_CONFIGS[wc.run_id]["env_id"])[0], - fdb_configured=lambda wc: bool( - _get_fdb_uenv_for_env(RUN_CONFIGS[wc.run_id]["env_id"])[0] - ), + fdb_configured=lambda wc: RUN_CONFIGS[wc.run_id]["output"] == "fdb", fdb_root=lambda wc: _resolve_fdb_root_for_run(wc.run_id), fdb_schema=str(Path("resources/fdb/realtime-varda.schema").resolve()), # fmt: off @@ -629,3 +626,32 @@ FDBEOF ) >{log} 2>&1 && touch {output.okfile} """ # fmt: on + + +# Fail fast: a run declaring `output: fdb` is relying on infrastructure and on its +# own referenced inference config to actually deliver FDB output -- catch the two +# ways that silently doesn't happen, at workflow-build time, before any (expensive) +# inference job runs: +# - no FDB uenv configured -> inference_check_fdb/inference_execute's own +# fdb_configured gate would just read False and FDB would never be touched at all, +# despite the declared intent. +# - the referenced inference config has zero `fdb:` output.tee blocks -> +# inference_check_fdb.py's completeness check is vacuously true when +# expected_blocks == 0 (0 complete blocks found >= 0 expected), reporting a false +# "hit" against whatever (possibly unrelated, possibly empty) data already sits in +# that FDB root. +for _run_id, _run_cfg in RUN_CONFIGS.items(): + if _run_cfg.get("output") != "fdb": + continue + if not _get_fdb_uenv_for_env(_run_cfg["env_id"])[0]: + raise ValueError( + f"Run '{_run_id}' declares output: fdb but profile.fdb.srun_uenv is not " + "set -- FDB archival/reads need a mounted FDB uenv to supply the native " + "libfdb5 library. Set profile.fdb.srun_uenv (e.g. 'fdb/5.21:v2')." + ) + if _count_fdb_blocks(_run_id) == 0: + raise ValueError( + f"Run '{_run_id}' declares output: fdb but its inference config " + f"({_run_cfg['config']}) has no 'fdb:' blocks under output.tee -- either " + "add one, or set output: grib to match what it actually writes." + ) diff --git a/workflow/rules/plot.smk b/workflow/rules/plot.smk index a2b3cf0a..a4c15835 100644 --- a/workflow/rules/plot.smk +++ b/workflow/rules/plot.smk @@ -64,18 +64,10 @@ rule plot_meteogram: if config.get("lapse_rate_correction", True) else "" ), - fdb_configured=lambda wc: bool( - _get_fdb_uenv_for_env(RUN_CONFIGS[wc.run_id]["env_id"])[0] - ), + fdb_configured=lambda wc: RUN_CONFIGS[wc.run_id]["output"] == "fdb", fdb_forecast=lambda wc: ( - "fdb:" - + ( - _get_fdb_roots(wc.run_id)[0] - if get_resource(wc, "write_to_global_fdb", False) - and _get_fdb_roots(wc.run_id)[0] - else _get_fdb_roots(wc.run_id)[1] - ) - if _get_fdb_uenv_for_env(RUN_CONFIGS[wc.run_id]["env_id"])[0] + "fdb:" + _resolve_fdb_root_for_run(wc.run_id) + if RUN_CONFIGS[wc.run_id]["output"] == "fdb" else "" ), shell: @@ -86,9 +78,9 @@ rule plot_meteogram: BASELINE_STEPS=({params.baseline_steps:q}) BASELINE_LABELS=({params.baseline_labels:q}) - # If this run's environment archives to FDB, its forecast output lives - # there (never in local grib -- see workflow/rules/inference.smk), so - # always prefer the FDB source over the (nonexistent) local grib dir. + # If this run declares output: fdb, its forecast output lives there + # (never in local grib -- see workflow/rules/inference.smk), so always + # prefer the FDB source over the (nonexistent) local grib dir. # Reading from FDB needs pyfdb (an evalml dependency, see pyproject.toml) # and eccodes-cosmo-mars's MARS-namespace GRIB concepts (a repo-level # clone, resources/fdb/patch_metkit_language.py's docstring has details) @@ -157,18 +149,10 @@ rule plot_forecast_frame: (Path(OUT_ROOT) / f"data/runs/{wc.run_id}/{wc.init_time}/frames").resolve() ), accu=lambda wc: int(RUN_CONFIGS[wc.run_id]["steps"].split("/")[2]), - fdb_configured=lambda wc: bool( - _get_fdb_uenv_for_env(RUN_CONFIGS[wc.run_id]["env_id"])[0] - ), + fdb_configured=lambda wc: RUN_CONFIGS[wc.run_id]["output"] == "fdb", fdb_forecast=lambda wc: ( - "fdb:" - + ( - _get_fdb_roots(wc.run_id)[0] - if get_resource(wc, "write_to_global_fdb", False) - and _get_fdb_roots(wc.run_id)[0] - else _get_fdb_roots(wc.run_id)[1] - ) - if _get_fdb_uenv_for_env(RUN_CONFIGS[wc.run_id]["env_id"])[0] + "fdb:" + _resolve_fdb_root_for_run(wc.run_id) + if RUN_CONFIGS[wc.run_id]["output"] == "fdb" else "" ), shell: @@ -176,7 +160,7 @@ rule plot_forecast_frame: set -euo pipefail # See plot_meteogram for why this always wins over the (nonexistent) - # local grib dir when the run's environment archives to FDB, and why no + # local grib dir when the run declares output: fdb, and why no # squashfs-mount is needed to read it (pyfdb bundles its own native libs). INPUT_SRC={params.grib_out_dir:q} if [ "{params.fdb_configured}" = "True" ]; then diff --git a/workflow/rules/verification.smk b/workflow/rules/verification.smk index 43e267de..9579ff49 100644 --- a/workflow/rules/verification.smk +++ b/workflow/rules/verification.smk @@ -97,18 +97,10 @@ rule verification_metrics: if config.get("lapse_rate_correction", True) else "" ), - fdb_configured=lambda wc: bool( - _get_fdb_uenv_for_env(RUN_CONFIGS[wc.run_id]["env_id"])[0] - ), + fdb_configured=lambda wc: RUN_CONFIGS[wc.run_id]["output"] == "fdb", fdb_forecast=lambda wc: ( - "fdb:" - + ( - _get_fdb_roots(wc.run_id)[0] - if get_resource(wc, "write_to_global_fdb", False) - and _get_fdb_roots(wc.run_id)[0] - else _get_fdb_roots(wc.run_id)[1] - ) - if _get_fdb_uenv_for_env(RUN_CONFIGS[wc.run_id]["env_id"])[0] + "fdb:" + _resolve_fdb_root_for_run(wc.run_id) + if RUN_CONFIGS[wc.run_id]["output"] == "fdb" else "" ), shell: @@ -116,9 +108,9 @@ rule verification_metrics: set -euo pipefail # See plot_meteogram (plot.smk) for why this always wins over the - # (nonexistent) local grib dir when the run's environment archives to - # FDB, and why no squashfs-mount is needed to read it (pyfdb bundles - # its own native libs). + # (nonexistent) local grib dir when the run declares output: fdb, and + # why no squashfs-mount is needed to read it (pyfdb bundles its own + # native libs). FORECAST_SRC={params.grib_out_dir:q} if [ "{params.fdb_configured}" = "True" ]; then FORECAST_SRC={params.fdb_forecast:q} @@ -254,22 +246,20 @@ rule verification_scoremaps: truth_label=(config.get("truth") or {}).get("label", ""), reftimes=[t.strftime("%Y%m%d%H%M") for t in REFTIMES], run_root=lambda wc: (Path(OUT_ROOT) / f"data/runs/{wc.run_id}").resolve(), - fdb_configured=lambda wc: bool( - _get_fdb_uenv_for_env(RUN_CONFIGS[wc.run_id]["env_id"])[0] - ), + fdb_configured=lambda wc: RUN_CONFIGS[wc.run_id]["output"] == "fdb", # Bare root (no "fdb:" marker prefix): --fdb_root is consumed by - # iter_init_dirs/fdb_has_data, which add the marker themselves when - # constructing load_forecast_data's root argument for a discovered - # FDB-resident reftime -- unlike plot.smk/verification_metrics' single - # --forecast argument, which IS load_forecast_data's root directly. + # iter_init_dirs, which adds the marker itself when constructing + # load_forecast_data's root argument for each reftime -- unlike + # plot.smk/verification_metrics' single --forecast argument, which IS + # load_forecast_data's root directly. A run's declared `output` is the + # same for every reftime under its run_root (it's a property of the run + # config, not of any one reftime), so this is the one root iter_init_dirs + # uses for all of them -- it never needs to guess per-directory. fdb_root=lambda wc: ( - _get_fdb_roots(wc.run_id)[0] - if get_resource(wc, "write_to_global_fdb", False) - and _get_fdb_roots(wc.run_id)[0] - else _get_fdb_roots(wc.run_id)[1] - ) - if _get_fdb_uenv_for_env(RUN_CONFIGS[wc.run_id]["env_id"])[0] - else "", + _resolve_fdb_root_for_run(wc.run_id) + if RUN_CONFIGS[wc.run_id]["output"] == "fdb" + else "" + ), shell: """ set -euo pipefail diff --git a/workflow/scripts/verification_scoremaps.py b/workflow/scripts/verification_scoremaps.py index e17febcc..5b77820f 100644 --- a/workflow/scripts/verification_scoremaps.py +++ b/workflow/scripts/verification_scoremaps.py @@ -39,7 +39,6 @@ import xarray as xr from data_input import ( - fdb_has_data, load_forecast_data, load_truth_data, open_truth_zarr, @@ -79,21 +78,27 @@ def _season_of(dt: datetime) -> str: def iter_init_dirs( run_root: Path, fdb_root: str | None = None ) -> list[tuple[datetime, Path | str]]: - """Return ``(reftime, grib_dir)`` pairs for every complete init time. + """Return ``(reftime, source)`` pairs for every init-time directory found. Expects subdirectories named ``YYYYMMDDHHMI`` directly under *run_root*. - GRIB files may live either directly in the init-time directory or inside a - ``grib/`` subdirectory. - - If `fdb_root` is given, init-time directories with no local GRIB output - (e.g. runs that archive directly to FDB and never write local grib/ at - all) are additionally checked against it and returned as an ``"fdb:"`` - marker string in place of a real path -- see data_input.load_forecast_data's - routing on that prefix, which every caller of this function's output - already goes through. Presence is checked with fdb_has_data(), a bare - existence check: completeness was already established upstream, by the - Snakemake dependency on this reftime's inference_check_fdb/inference_execute - okfile, before this script ever runs. + ``fdb_root`` being set or ``None`` is the caller's explicit declaration of + this run's backend (derived from its config's ``output: fdb``/``grib`` + field -- see ``workflow/rules/verification.smk``'s ``fdb_root`` param), + not a hint to probe with: every directory found is assigned a source the + same way, with no per-directory existence check of either backend. + - ``output: fdb`` (``fdb_root`` given): every reftime maps to an + ``"fdb:"`` marker string -- see ``data_input.load_forecast_data``'s + routing on that prefix, which every caller of this function's output + already goes through. + - ``output: grib`` (``fdb_root`` is ``None``): every reftime maps to its + GRIB directory (files may live directly in the init-time directory or + inside a ``grib/`` subdirectory -- a layout detail, not a backend guess). + + Completeness is the caller's job, not this function's: ``main()`` requires + every configured reftime to be discovered here, raising a hard error + otherwise (see its module docstring) -- consistent with the Snakemake + dependency on this reftime's inference_check_fdb/inference_execute okfile + already having established completeness before this script ever runs. """ result = [] for d in sorted(run_root.iterdir()): @@ -103,13 +108,11 @@ def iter_init_dirs( reftime = datetime.strptime(d.name, DATETIME_FMT) except ValueError: continue - grib_dir = d / "grib" if (d / "grib").is_dir() else d - if any(grib_dir.glob("*.grib")): - result.append((reftime, grib_dir)) - elif fdb_root is not None and fdb_has_data(fdb_root, reftime): + if fdb_root is not None: result.append((reftime, f"fdb:{fdb_root}")) else: - LOG.debug("No GRIB files in %s, skipping", grib_dir) + grib_dir = d / "grib" if (d / "grib").is_dir() else d + result.append((reftime, grib_dir)) return result diff --git a/workflow/tools/config.schema.json b/workflow/tools/config.schema.json index 7f85fa37..84db9ce3 100644 --- a/workflow/tools/config.schema.json +++ b/workflow/tools/config.schema.json @@ -469,6 +469,16 @@ "title": "Disable Local Eccodes Definitions", "type": "boolean" }, + "output": { + "default": "grib", + "description": "Where this run's forecast output is stored, and therefore where downstream reads (verification, meteograms, frames, scoremaps) should look for it. Must match what the referenced inference `config`'s own output.tee actually writes to (an 'fdb' tee block requires output: fdb here, and vice versa) -- this is the single declared source of truth for backend routing; nothing probes for data presence or parses the inference config to guess it.", + "enum": [ + "grib", + "fdb" + ], + "title": "Output", + "type": "string" + }, "config": { "anyOf": [ { @@ -1005,6 +1015,16 @@ "title": "Disable Local Eccodes Definitions", "type": "boolean" }, + "output": { + "default": "grib", + "description": "Where this run's forecast output is stored, and therefore where downstream reads (verification, meteograms, frames, scoremaps) should look for it. Must match what the referenced inference `config`'s own output.tee actually writes to (an 'fdb' tee block requires output: fdb here, and vice versa) -- this is the single declared source of truth for backend routing; nothing probes for data presence or parses the inference config to guess it.", + "enum": [ + "grib", + "fdb" + ], + "title": "Output", + "type": "string" + }, "config": { "anyOf": [ { From d2d33deacfcb1f66a510f69f44ca1af24244f15a Mon Sep 17 00:00:00 2001 From: Carlos Osuna Date: Wed, 21 Oct 2026 09:47:03 +0200 Subject: [PATCH 14/21] add a root key --- config/varda-single-1.0-inference-only.yaml | 4 -- src/evalml/config.py | 32 ++++++----- workflow/rules/inference.smk | 63 ++++++--------------- workflow/scripts/inference_check_fdb.py | 9 +-- workflow/tools/config.schema.json | 45 ++++++++------- 5 files changed, 67 insertions(+), 86 deletions(-) diff --git a/config/varda-single-1.0-inference-only.yaml b/config/varda-single-1.0-inference-only.yaml index 91959060..ceabcd56 100644 --- a/config/varda-single-1.0-inference-only.yaml +++ b/config/varda-single-1.0-inference-only.yaml @@ -33,10 +33,6 @@ profile: fdb: srun_uenv: "fdb/5.19:v2" srun_view: "realtime" - # Base dir for the shared store; actual root per checkpoint is - # {fdb_root_global}/{env_id}. TODO: point at the real shared location. - fdb_root_global: "/scratch/mch/cosuna/evalml_fdb/global-fdb" - write_to_global_fdb: false executor: slurm global_resources: gpus: 16 diff --git a/src/evalml/config.py b/src/evalml/config.py index 69ebf18c..af7f7abe 100644 --- a/src/evalml/config.py +++ b/src/evalml/config.py @@ -79,20 +79,6 @@ class InferenceResources(BaseModel): None, description="FDB uenv view name to use when bundling FDB libraries, e.g. 'realtime'.", ) - fdb_root_global: str | None = Field( - None, - description="Base directory for the shared/global FDB store. Each checkpoint gets " - "its own dedicated root at '{fdb_root_global}/{env_id}', keyed by the env_id hash " - "(derived from the checkpoint) so different checkpoints never mix data in the same " - "FDB instance. When set, inference is skipped for a given (run_id, init_time) if " - "this root already holds every requested lead time.", - ) - write_to_global_fdb: bool = Field( - False, - description="If true, inference writes directly into " - "'{fdb_root_global}/{env_id}' instead of the local per-checkpoint root. " - "Requires fdb_root_global to be set.", - ) class RunConfig(BaseModel): @@ -138,10 +124,28 @@ class RunConfig(BaseModel): "is the single declared source of truth for backend routing; nothing probes for " "data presence or parses the inference config to guess it.", ) + root: str | None = Field( + None, + description="Only meaningful when output: fdb. Root directory of a shared/global " + "FDB store this run reads from and writes into, used as-is (same pattern as " + "baseline.root) -- the caller is responsible for pointing different checkpoints " + "at different roots if they shouldn't share one. When unset (the default), this " + "run's FDB output lives in a local per-checkpoint directory under the output root.", + ) config: Dict[str, Any] | str model_config = {"extra": "forbid"} + @model_validator(mode="after") + def _root_requires_fdb_output(self) -> "RunConfig": + if self.root and self.output != "fdb": + raise ValueError( + "root is set but output is not 'fdb' -- a custom root only applies to " + "FDB-backed runs (grib output always lives under the local run " + "directory); set output: fdb or remove root." + ) + return self + @field_validator("steps") def validate_steps(cls, v: str) -> str: if "/" not in v: diff --git a/workflow/rules/inference.smk b/workflow/rules/inference.smk index 28b56081..971c2362 100644 --- a/workflow/rules/inference.smk +++ b/workflow/rules/inference.smk @@ -210,29 +210,6 @@ def _get_fdb_uenv_for_env(env_id): return uenv, view -def _get_fdb_roots(run_id): - """Return (global_root, local_root) for a run, both keyed by env_id. - - Each checkpoint (env_id) gets its own dedicated FDB instance, so different - checkpoints never mix data in the same root. global_root is '' if - profile.fdb.fdb_root_global is not configured. - """ - env_id = RUN_CONFIGS[run_id]["env_id"] - global_base = get_resource_by_env(env_id, "fdb_root_global", "") - global_root = f"{global_base}/{env_id}" if global_base else "" - local_root = str((OUT_ROOT / f"data/fdb/{env_id}").resolve()) - return global_root, local_root - - -def get_resource_by_env(env_id, field: str, default): - """Like get_resource, but usable outside a rule (no wildcards available).""" - fdb = (config.get("profile") or {}).get("fdb") or {} - if isinstance(fdb, dict): - return fdb.get(field, default) or default - else: - return getattr(fdb, field, default) or default - - def _count_fdb_blocks(run_id): """Count distinct 'fdb:' targets in this run's own inference config output.tee. @@ -250,10 +227,11 @@ def _count_fdb_blocks(run_id): checkpoint inference_check_fdb: """Check whether (run_id, init_time) output already exists, complete, in FDB. - Checks the global root first (if profile.fdb.fdb_root_global is set), then the - local per-checkpoint root, in that order. A 'hit' requires every requested lead - time to be present in one root -- see inference_check_fdb.py for why no further - MARS-key disambiguation (class/expver/model) is needed. Writes 'hit\\t' or + Checks this run's one configured FDB root (see _resolve_fdb_root_for_run -- + the run's own `root:` field if set, otherwise its local per-checkpoint + directory). A 'hit' requires every requested lead time to be present in + that root -- see inference_check_fdb.py for why no further MARS-key + disambiguation (class/expver/model) is needed. Writes 'hit\\t' or 'miss\\t' to the status file. Runs inside the squashfs venv so that pyfdb and the native FDB C library are available. Skipped entirely (writes 'miss') unless this run's own config declares `output: fdb` (see the fail-fast validation below this @@ -272,8 +250,7 @@ checkpoint inference_check_fdb: image_path=lambda wc, input: str(Path(input.image).resolve()), fdb_configured=lambda wc: RUN_CONFIGS[wc.run_id]["output"] == "fdb", fdb_uenv=lambda wc: _get_fdb_uenv_for_env(RUN_CONFIGS[wc.run_id]["env_id"])[0], - fdb_root_global=lambda wc: _get_fdb_roots(wc.run_id)[0], - fdb_root_local=lambda wc: _get_fdb_roots(wc.run_id)[1], + fdb_root=lambda wc: _resolve_fdb_root_for_run(wc.run_id), fdb_schema=str(Path("resources/fdb/realtime-varda.schema").resolve()), fdb_view=lambda wc: get_resource(wc, "srun_view", "realtime") or "realtime", date=lambda wc: datetime.strptime(wc.init_time, "%Y%m%d%H%M").strftime("%Y%m%d"), @@ -288,10 +265,6 @@ checkpoint inference_check_fdb: printf 'miss\t' > {output.status} echo "FDB check skipped: FDB not configured for this environment" else - GLOBAL_ARG="" - if [ -n "{params.fdb_root_global}" ]; then - GLOBAL_ARG="--fdb-root {params.fdb_root_global}" - fi # Mount the venv (pyfdb, python) and the FDB uenv (native libs, config, # at its required /user-environment path) side by side in one call. # squashfs-mount needs the mountpoint dir to already exist -- /user-environment @@ -303,8 +276,8 @@ checkpoint inference_check_fdb: squashfs-mount -s "{params.image_path}:$VENV_MOUNT,$FDB_SQUASHFS:/user-environment" -- bash -c ' source '"$VENV_MOUNT"'/bin/activate export LD_LIBRARY_PATH=/user-environment/env/{params.fdb_view}/lib64:/user-environment/env/{params.fdb_view}/lib${{LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}} - python {input.script} '"$GLOBAL_ARG"' \ - --fdb-root {params.fdb_root_local} \ + python {input.script} \ + --fdb-root {params.fdb_root} \ --fdb-schema {params.fdb_schema} \ --date {params.date} \ --time {params.time} \ @@ -318,15 +291,15 @@ checkpoint inference_check_fdb: def _resolve_fdb_root_for_run(run_id): - """Which FDB root this run's own inference_execute would write fresh output - into: the global one if write_to_global_fdb is set (and configured), - otherwise the local per-checkpoint root. Same logic as inference_execute's - own `fdb_root` param -- factored out so the read side can agree with it.""" - global_root, local_root = _get_fdb_roots(run_id) + """This run's one FDB root: its own `root:` config value if set (a + shared/global store, used as-is -- same pattern as baseline.root), + otherwise a local directory dedicated to this checkpoint (env_id-keyed) + under the output root.""" + custom_root = RUN_CONFIGS[run_id].get("root") + if custom_root: + return custom_root env_id = RUN_CONFIGS[run_id]["env_id"] - if get_resource_by_env(env_id, "write_to_global_fdb", False) and global_root: - return global_root - return local_root + return str((OUT_ROOT / f"data/fdb/{env_id}").resolve()) def _get_inference_status(wc): @@ -499,8 +472,8 @@ rule inference_execute: cd {params.workdir} - # Write into this checkpoint's dedicated FDB root (env_id-keyed): the global - # one if write_to_global_fdb is set, otherwise the local per-checkpoint root. + # Write into this run's one FDB root (see _resolve_fdb_root_for_run): its + # own `root:` config value if set, otherwise its local per-checkpoint root. # Exported so _run_inference (executed by squashfs-mount in a fresh bash) can see it. export FDB5_CONFIG_OVERRIDE="" if [ "{params.fdb_configured}" = "True" ]; then diff --git a/workflow/scripts/inference_check_fdb.py b/workflow/scripts/inference_check_fdb.py index 977ee639..b1d45ca7 100644 --- a/workflow/scripts/inference_check_fdb.py +++ b/workflow/scripts/inference_check_fdb.py @@ -1,10 +1,11 @@ """Check whether inference output for a given init_time already exists in FDB. Each FDB root passed on the command line is assumed dedicated to a single -checkpoint (the caller is responsible for pointing --fdb-root at an -env_id-keyed directory, e.g. '{fdb_root_global}/{env_id}' or -'{local_base}/{env_id}'), so no cross-checkpoint MARS-key disambiguation -(class/expver/model) is needed here. +checkpoint (the caller passes this run's own resolved root -- see +`_resolve_fdb_root_for_run` in workflow/rules/inference.smk: either the run's +explicit `root:` config value, or else its local per-checkpoint directory), +so no cross-checkpoint MARS-key disambiguation (class/expver/model) is +needed here. A single init_time's output can still be split into several distinct *blocks* within that one root -- e.g. an ICON-grid target and an IFS-grid diff --git a/workflow/tools/config.schema.json b/workflow/tools/config.schema.json index 84db9ce3..a7dbaed4 100644 --- a/workflow/tools/config.schema.json +++ b/workflow/tools/config.schema.json @@ -479,6 +479,19 @@ "title": "Output", "type": "string" }, + "root": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Only meaningful when output: fdb. Root directory of a shared/global FDB store this run reads from and writes into, used as-is (same pattern as baseline.root) -- the caller is responsible for pointing different checkpoints at different roots if they shouldn't share one. When unset (the default), this run's FDB output lives in a local per-checkpoint directory under the output root.", + "title": "Root" + }, "config": { "anyOf": [ { @@ -646,25 +659,6 @@ "default": null, "description": "FDB uenv view name to use when bundling FDB libraries, e.g. 'realtime'.", "title": "Srun View" - }, - "fdb_root_global": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "default": null, - "description": "Base directory for the shared/global FDB store. Each checkpoint gets its own dedicated root at '{fdb_root_global}/{env_id}', keyed by the env_id hash (derived from the checkpoint) so different checkpoints never mix data in the same FDB instance. When set, inference is skipped for a given (run_id, init_time) if this root already holds every requested lead time.", - "title": "Fdb Root Global" - }, - "write_to_global_fdb": { - "default": false, - "description": "If true, inference writes directly into '{fdb_root_global}/{env_id}' instead of the local per-checkpoint root. Requires fdb_root_global to be set.", - "title": "Write To Global Fdb", - "type": "boolean" } }, "title": "InferenceResources", @@ -1025,6 +1019,19 @@ "title": "Output", "type": "string" }, + "root": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Only meaningful when output: fdb. Root directory of a shared/global FDB store this run reads from and writes into, used as-is (same pattern as baseline.root) -- the caller is responsible for pointing different checkpoints at different roots if they shouldn't share one. When unset (the default), this run's FDB output lives in a local per-checkpoint directory under the output root.", + "title": "Root" + }, "config": { "anyOf": [ { From c01053ec79ca8c48b14335bab0f6cd7088a56516 Mon Sep 17 00:00:00 2001 From: Carlos Osuna Date: Wed, 21 Oct 2026 11:15:55 +0200 Subject: [PATCH 15/21] adding CLCx --- ...aler-global_trimedge_multi-fdb-parity.yaml | 140 +++++++++++++----- ...-downscaler-global_trimedge_multi-fdb.yaml | 140 +++++++++++++----- 2 files changed, 198 insertions(+), 82 deletions(-) diff --git a/resources/inference/configs/sgm-temporal-downscaler-global_trimedge_multi-fdb-parity.yaml b/resources/inference/configs/sgm-temporal-downscaler-global_trimedge_multi-fdb-parity.yaml index a48d5dfb..b881e4df 100644 --- a/resources/inference/configs/sgm-temporal-downscaler-global_trimedge_multi-fdb-parity.yaml +++ b/resources/inference/configs/sgm-temporal-downscaler-global_trimedge_multi-fdb-parity.yaml @@ -27,6 +27,47 @@ input: scale: 0.001 # convert from kg m-2 to m offset: 0 param: TOT_PREC + # filters see the model name (tcc) with anemoi-inference <= 0.11.1 and the GRIB name (CLCT) from 0.11.2 + - forward_transform_filter: # cloud cover percent -> fraction, as in training + rescale: + scale: 0.01 + offset: 0 + param: tcc + - forward_transform_filter: # cloud cover percent -> fraction, as in training + rescale: + scale: 0.01 + offset: 0 + param: CLCT + - forward_transform_filter: # cloud cover percent -> fraction, as in training + rescale: + scale: 0.01 + offset: 0 + param: lcc + - forward_transform_filter: # cloud cover percent -> fraction, as in training + rescale: + scale: 0.01 + offset: 0 + param: CLCL + - forward_transform_filter: # cloud cover percent -> fraction, as in training + rescale: + scale: 0.01 + offset: 0 + param: mcc + - forward_transform_filter: # cloud cover percent -> fraction, as in training + rescale: + scale: 0.01 + offset: 0 + param: CLCM + - forward_transform_filter: # cloud cover percent -> fraction, as in training + rescale: + scale: 0.01 + offset: 0 + param: hcc + - forward_transform_filter: # cloud cover percent -> fraction, as in training + rescale: + scale: 0.01 + offset: 0 + param: CLCH namer: &namer rules: # ICON short names (uppercase) @@ -52,8 +93,16 @@ input: - 2t - - shortName: TD_2M - 2d - - - shortName: T_G + - - shortName: SKT - skt + - - shortName: CLCT + - tcc + - - shortName: CLCL + - lcc + - - shortName: CLCM + - mcc + - - shortName: CLCH + - hcc - - shortName: U_10M - 10u - - shortName: V_10M @@ -62,46 +111,6 @@ input: - lsm - - shortName: TOT_PREC - tp - # IFS short names (lowercase) for the global grid output - - - shortName: t - levtype: pl - - t_{level} - - - shortName: u - levtype: pl - - u_{level} - - - shortName: v - levtype: pl - - v_{level} - - - shortName: w - levtype: pl - - w_{level} - - - shortName: q - levtype: pl - - q_{level} - - - shortName: z - levtype: pl - - z_{level} - - - shortName: z - levtype: sfc - - z - - - shortName: 10u - - 10u - - - shortName: 10v - - 10v - - - shortName: 2t - - 2t - - - shortName: 2d - - 2d - - - shortName: msl - - msl - - - shortName: sp - - sp - - - shortName: lsm - - lsm - - - shortName: skt - - skt - - - shortName: tp - - tp - global: grib: path: forecaster/ifs* @@ -140,6 +149,47 @@ output: - extract_mask: # removes global points mask: "lam_0/cutout_mask" as_slice: true + # filters see the model name (tcc) with anemoi-inference <= 0.11.1 and the GRIB name (CLCT) from 0.11.2 + - forward_transform_filter: # cloud cover fraction -> percent, before copy-prognostic + rescale: + scale: 100 + offset: 0 + param: tcc + - forward_transform_filter: # cloud cover fraction -> percent, before copy-prognostic + rescale: + scale: 100 + offset: 0 + param: CLCT + - forward_transform_filter: # cloud cover fraction -> percent, before copy-prognostic + rescale: + scale: 100 + offset: 0 + param: lcc + - forward_transform_filter: # cloud cover fraction -> percent, before copy-prognostic + rescale: + scale: 100 + offset: 0 + param: CLCL + - forward_transform_filter: # cloud cover fraction -> percent, before copy-prognostic + rescale: + scale: 100 + offset: 0 + param: mcc + - forward_transform_filter: # cloud cover fraction -> percent, before copy-prognostic + rescale: + scale: 100 + offset: 0 + param: CLCM + - forward_transform_filter: # cloud cover fraction -> percent, before copy-prognostic + rescale: + scale: 100 + offset: 0 + param: hcc + - forward_transform_filter: # cloud cover fraction -> percent, before copy-prognostic + rescale: + scale: 100 + offset: 0 + param: CLCH # here, the trimedge mask can be specified when available # at overlap steps (multiples of the forecaster stride) replace the # re-predicted prognostics with the forecaster's boundary values @@ -199,6 +249,14 @@ output: slor: {"param": 163, "shortName": "slor"} SSO_STDH: {"param": 160, "shortName": "sdor"} sdor: {"param": 160, "shortName": "sdor"} + CLCT: {"param": 164, "shortName": "tcc"} + tcc: {"param": 164, "shortName": "tcc"} + CLCL: {"param": 186, "shortName": "lcc"} + lcc: {"param": 186, "shortName": "lcc"} + CLCM: {"param": 187, "shortName": "mcc"} + mcc: {"param": 187, "shortName": "mcc"} + CLCH: {"param": 188, "shortName": "hcc"} + hcc: {"param": 188, "shortName": "hcc"} TOT_PREC: {"param": 228, "shortName": "tp"} tp: {"param": 228, "shortName": "tp"} z: {"param": 129, "shortName": "z"} diff --git a/resources/inference/configs/sgm-temporal-downscaler-global_trimedge_multi-fdb.yaml b/resources/inference/configs/sgm-temporal-downscaler-global_trimedge_multi-fdb.yaml index a288d8c5..33b4eb30 100644 --- a/resources/inference/configs/sgm-temporal-downscaler-global_trimedge_multi-fdb.yaml +++ b/resources/inference/configs/sgm-temporal-downscaler-global_trimedge_multi-fdb.yaml @@ -17,6 +17,47 @@ input: scale: 0.001 # convert from kg m-2 to m offset: 0 param: TOT_PREC + # filters see the model name (tcc) with anemoi-inference <= 0.11.1 and the GRIB name (CLCT) from 0.11.2 + - forward_transform_filter: # cloud cover percent -> fraction, as in training + rescale: + scale: 0.01 + offset: 0 + param: tcc + - forward_transform_filter: # cloud cover percent -> fraction, as in training + rescale: + scale: 0.01 + offset: 0 + param: CLCT + - forward_transform_filter: # cloud cover percent -> fraction, as in training + rescale: + scale: 0.01 + offset: 0 + param: lcc + - forward_transform_filter: # cloud cover percent -> fraction, as in training + rescale: + scale: 0.01 + offset: 0 + param: CLCL + - forward_transform_filter: # cloud cover percent -> fraction, as in training + rescale: + scale: 0.01 + offset: 0 + param: mcc + - forward_transform_filter: # cloud cover percent -> fraction, as in training + rescale: + scale: 0.01 + offset: 0 + param: CLCM + - forward_transform_filter: # cloud cover percent -> fraction, as in training + rescale: + scale: 0.01 + offset: 0 + param: hcc + - forward_transform_filter: # cloud cover percent -> fraction, as in training + rescale: + scale: 0.01 + offset: 0 + param: CLCH namer: &namer rules: # ICON short names (uppercase) @@ -42,8 +83,16 @@ input: - 2t - - shortName: TD_2M - 2d - - - shortName: T_G + - - shortName: SKT - skt + - - shortName: CLCT + - tcc + - - shortName: CLCL + - lcc + - - shortName: CLCM + - mcc + - - shortName: CLCH + - hcc - - shortName: U_10M - 10u - - shortName: V_10M @@ -52,46 +101,6 @@ input: - lsm - - shortName: TOT_PREC - tp - # IFS short names (lowercase) for the global grid output - - - shortName: t - levtype: pl - - t_{level} - - - shortName: u - levtype: pl - - u_{level} - - - shortName: v - levtype: pl - - v_{level} - - - shortName: w - levtype: pl - - w_{level} - - - shortName: q - levtype: pl - - q_{level} - - - shortName: z - levtype: pl - - z_{level} - - - shortName: z - levtype: sfc - - z - - - shortName: 10u - - 10u - - - shortName: 10v - - 10v - - - shortName: 2t - - 2t - - - shortName: 2d - - 2d - - - shortName: msl - - msl - - - shortName: sp - - sp - - - shortName: lsm - - lsm - - - shortName: skt - - skt - - - shortName: tp - - tp - global: grib: path: forecaster/ifs* @@ -131,6 +140,47 @@ output: - extract_mask: # removes global points mask: "lam_0/cutout_mask" as_slice: true + # filters see the model name (tcc) with anemoi-inference <= 0.11.1 and the GRIB name (CLCT) from 0.11.2 + - forward_transform_filter: # cloud cover fraction -> percent, before copy-prognostic + rescale: + scale: 100 + offset: 0 + param: tcc + - forward_transform_filter: # cloud cover fraction -> percent, before copy-prognostic + rescale: + scale: 100 + offset: 0 + param: CLCT + - forward_transform_filter: # cloud cover fraction -> percent, before copy-prognostic + rescale: + scale: 100 + offset: 0 + param: lcc + - forward_transform_filter: # cloud cover fraction -> percent, before copy-prognostic + rescale: + scale: 100 + offset: 0 + param: CLCL + - forward_transform_filter: # cloud cover fraction -> percent, before copy-prognostic + rescale: + scale: 100 + offset: 0 + param: mcc + - forward_transform_filter: # cloud cover fraction -> percent, before copy-prognostic + rescale: + scale: 100 + offset: 0 + param: CLCM + - forward_transform_filter: # cloud cover fraction -> percent, before copy-prognostic + rescale: + scale: 100 + offset: 0 + param: hcc + - forward_transform_filter: # cloud cover fraction -> percent, before copy-prognostic + rescale: + scale: 100 + offset: 0 + param: CLCH - fdb: encoding: @@ -175,6 +225,14 @@ output: slor: {"param": 163, "shortName": "slor"} SSO_STDH: {"param": 160, "shortName": "sdor"} sdor: {"param": 160, "shortName": "sdor"} + CLCT: {"param": 164, "shortName": "tcc"} + tcc: {"param": 164, "shortName": "tcc"} + CLCL: {"param": 186, "shortName": "lcc"} + lcc: {"param": 186, "shortName": "lcc"} + CLCM: {"param": 187, "shortName": "mcc"} + mcc: {"param": 187, "shortName": "mcc"} + CLCH: {"param": 188, "shortName": "hcc"} + hcc: {"param": 188, "shortName": "hcc"} TOT_PREC: {"param": 228, "shortName": "tp"} tp: {"param": 228, "shortName": "tp"} z: {"param": 129, "shortName": "z"} From 86ad7e8d28a7163a0babb0d39197e8574fe4d396 Mon Sep 17 00:00:00 2001 From: Carlos Osuna Date: Wed, 21 Oct 2026 12:06:36 +0200 Subject: [PATCH 16/21] merge configs. --- ...aler-global_trimedge_multi-fdb-parity.yaml | 273 ------------------ ...-downscaler-global_trimedge_multi-fdb.yaml | 249 ---------------- 2 files changed, 522 deletions(-) delete mode 100644 resources/inference/configs/sgm-temporal-downscaler-global_trimedge_multi-fdb-parity.yaml delete mode 100644 resources/inference/configs/sgm-temporal-downscaler-global_trimedge_multi-fdb.yaml diff --git a/resources/inference/configs/sgm-temporal-downscaler-global_trimedge_multi-fdb-parity.yaml b/resources/inference/configs/sgm-temporal-downscaler-global_trimedge_multi-fdb-parity.yaml deleted file mode 100644 index b881e4df..00000000 --- a/resources/inference/configs/sgm-temporal-downscaler-global_trimedge_multi-fdb-parity.yaml +++ /dev/null @@ -1,273 +0,0 @@ -runner: temporal_downscaler - -# FDB counterpart to sgm-temporal-downscaler-global_trimedge_multi.yaml ("main"), -# built for a fair inference/downstream timing comparison between the two -# storage backends: same input/post_processors/copy-prognostic-from-forecaster -# filter as main (identical forecast computation), only output.tee differs -# (fdb: blocks instead of grib:, structured like -# sgm-temporal-downscaler-global_trimedge_multi-fdb.yaml's already-validated -# FDB encoding/modifiers). Requires anemoi-inference installed from -# /scratch/mch/cosuna/evalml_fdb/anemoi-inference (PR #573 fix + local custom -# FDBOutput class) -- see varda-single-1.0-fdb-timing.yaml. - -input: - cutout: - - lam_0: - grib: - path: forecaster/20* - pre_processors: - - forward_transform_filter: - # does not have an effect if the temporal downscaler does not use tp as prognostic variable - rescale: - scale: 0.001 # convert from kg m-2 to m - offset: 0 - param: tp - - forward_transform_filter: - rescale: - scale: 0.001 # convert from kg m-2 to m - offset: 0 - param: TOT_PREC - # filters see the model name (tcc) with anemoi-inference <= 0.11.1 and the GRIB name (CLCT) from 0.11.2 - - forward_transform_filter: # cloud cover percent -> fraction, as in training - rescale: - scale: 0.01 - offset: 0 - param: tcc - - forward_transform_filter: # cloud cover percent -> fraction, as in training - rescale: - scale: 0.01 - offset: 0 - param: CLCT - - forward_transform_filter: # cloud cover percent -> fraction, as in training - rescale: - scale: 0.01 - offset: 0 - param: lcc - - forward_transform_filter: # cloud cover percent -> fraction, as in training - rescale: - scale: 0.01 - offset: 0 - param: CLCL - - forward_transform_filter: # cloud cover percent -> fraction, as in training - rescale: - scale: 0.01 - offset: 0 - param: mcc - - forward_transform_filter: # cloud cover percent -> fraction, as in training - rescale: - scale: 0.01 - offset: 0 - param: CLCM - - forward_transform_filter: # cloud cover percent -> fraction, as in training - rescale: - scale: 0.01 - offset: 0 - param: hcc - - forward_transform_filter: # cloud cover percent -> fraction, as in training - rescale: - scale: 0.01 - offset: 0 - param: CLCH - namer: &namer - rules: - # ICON short names (uppercase) - - - shortName: T - - t_{level} - - - shortName: U - - u_{level} - - - shortName: V - - v_{level} - - - shortName: W - - w_{level} - - - shortName: QV - - q_{level} - - - shortName: FI - - z_{level} - - - shortName: PMSL - - msl - - - shortName: FIS - - z - - - shortName: PS - - sp - - - shortName: T_2M - - 2t - - - shortName: TD_2M - - 2d - - - shortName: SKT - - skt - - - shortName: CLCT - - tcc - - - shortName: CLCL - - lcc - - - shortName: CLCM - - mcc - - - shortName: CLCH - - hcc - - - shortName: U_10M - - 10u - - - shortName: V_10M - - 10v - - - shortName: FR_LAND - - lsm - - - shortName: TOT_PREC - - tp - - global: - grib: - path: forecaster/ifs* - namer: *namer - -constant_forcings: - test: - use_original_paths: true - -patch_metadata: resources/sgm-temporal-downscaler-ich1-oper-patch.yaml - -post_processors: - - accumulate_from_start_of_forecast: # accumulate tp from start of forecast - accumulations: - - tp - - forward_transform_filter: - rescale: - scale: 1000 # convert units from m to kg m-2 - offset: 0 - param: tp - - forward_transform_filter: - rescale: - scale: 1000 # convert units from m to kg m-2 - offset: 0 - param: TOT_PREC - -output: - tee: - - fdb: - encoding: - typeOfGeneratingProcess: 2 - generatingProcessIdentifier: 180 - templates: - samples: resources/templates_index_icon.yaml - post_processors: - - extract_mask: # removes global points - mask: "lam_0/cutout_mask" - as_slice: true - # filters see the model name (tcc) with anemoi-inference <= 0.11.1 and the GRIB name (CLCT) from 0.11.2 - - forward_transform_filter: # cloud cover fraction -> percent, before copy-prognostic - rescale: - scale: 100 - offset: 0 - param: tcc - - forward_transform_filter: # cloud cover fraction -> percent, before copy-prognostic - rescale: - scale: 100 - offset: 0 - param: CLCT - - forward_transform_filter: # cloud cover fraction -> percent, before copy-prognostic - rescale: - scale: 100 - offset: 0 - param: lcc - - forward_transform_filter: # cloud cover fraction -> percent, before copy-prognostic - rescale: - scale: 100 - offset: 0 - param: CLCL - - forward_transform_filter: # cloud cover fraction -> percent, before copy-prognostic - rescale: - scale: 100 - offset: 0 - param: mcc - - forward_transform_filter: # cloud cover fraction -> percent, before copy-prognostic - rescale: - scale: 100 - offset: 0 - param: CLCM - - forward_transform_filter: # cloud cover fraction -> percent, before copy-prognostic - rescale: - scale: 100 - offset: 0 - param: hcc - - forward_transform_filter: # cloud cover fraction -> percent, before copy-prognostic - rescale: - scale: 100 - offset: 0 - param: CLCH - # here, the trimedge mask can be specified when available - # at overlap steps (multiples of the forecaster stride) replace the - # re-predicted prognostics with the forecaster's boundary values - - forward_transform_filter: - copy-prognostic-from-forecaster: - forecaster_path: forecaster/20* - common_leadtime: 6h - params_to_keep: [tp] # NOTE, if more diagnostics are added, update accordingly. - namer: *namer - - - fdb: - encoding: - typeOfGeneratingProcess: 2 - generatingProcessIdentifier: 180 - # localNumberOfExperiment is not available in IFS/ECMWF GRIB templates; - # omitting it means marsClass=od (GPI=180 fallback) for the global grid - # Set native marsStream/marsType codetable keys so g2_mars_labeling returns - # stream=enfo and type=fc. Cannot use a concept override in local.98.def - # for marsStream because it conflicts with the codetable key that - # g2_mars_labeling expects as input. Setting them natively is portable - # across eccodes builds. - marsStream: enfo - marsType: fc - templates: - samples: resources/templates_index_ifs.yaml - post_processors: - - extract_mask: # removes lam points - mask: "lam_0/cutout_mask" - as_slice: true - inverse: true - - assign_mask: # fill local/global overlapping points with nan - mask: "global/cutout_mask" - - forward_transform_filter: - copy-prognostic-from-forecaster: - forecaster_path: forecaster/ifs* - common_leadtime: 6h - params_to_keep: [tp] # NOTE, if more diagnostics are added, update accordingly. - namer: *namer - modifiers: - - patches: - - variable: - U_10M: {"param": 165, "shortName": "10u"} - 10u: {"param": 165, "shortName": "10u"} - V_10M: {"param": 166, "shortName": "10v"} - 10v: {"param": 166, "shortName": "10v"} - TD_2M: {"param": 168, "shortName": "2d"} - 2d: {"param": 168, "shortName": "2d"} - T_2M: {"param": 167, "shortName": "2t"} - 2t: {"param": 167, "shortName": "2t"} - FR_LAND: {"param": 172, "shortName": "lsm"} - lsm: {"param": 172, "shortName": "lsm"} - PMSL: {"param": 151, "shortName": "msl"} - msl: {"param": 151, "shortName": "msl"} - PS: {"param": 134, "shortName": "sp"} - sp: {"param": 134, "shortName": "sp"} - SSO_SIGMA: {"param": 163, "shortName": "slor"} - slor: {"param": 163, "shortName": "slor"} - SSO_STDH: {"param": 160, "shortName": "sdor"} - sdor: {"param": 160, "shortName": "sdor"} - CLCT: {"param": 164, "shortName": "tcc"} - tcc: {"param": 164, "shortName": "tcc"} - CLCL: {"param": 186, "shortName": "lcc"} - lcc: {"param": 186, "shortName": "lcc"} - CLCM: {"param": 187, "shortName": "mcc"} - mcc: {"param": 187, "shortName": "mcc"} - CLCH: {"param": 188, "shortName": "hcc"} - hcc: {"param": 188, "shortName": "hcc"} - TOT_PREC: {"param": 228, "shortName": "tp"} - tp: {"param": 228, "shortName": "tp"} - z: {"param": 129, "shortName": "z"} - "^q_(\\d+)$": {"param": 133, "shortName": "q"} - "^t_(\\d+)$": {"param": 130, "shortName": "t"} - "^u_(\\d+)$": {"param": 131, "shortName": "u"} - "^v_(\\d+)$": {"param": 132, "shortName": "v"} - "^w_(\\d+)$": {"param": 135, "shortName": "w"} - "^z_(\\d+)$": {"param": 129, "shortName": "z"} - -# silenced due to bug in anemoi-inference for multi-step temporal downscalers, can be removed when fixed -verbosity: 0 -allow_nans: true -output_frequency: "1h" diff --git a/resources/inference/configs/sgm-temporal-downscaler-global_trimedge_multi-fdb.yaml b/resources/inference/configs/sgm-temporal-downscaler-global_trimedge_multi-fdb.yaml deleted file mode 100644 index 33b4eb30..00000000 --- a/resources/inference/configs/sgm-temporal-downscaler-global_trimedge_multi-fdb.yaml +++ /dev/null @@ -1,249 +0,0 @@ -runner: temporal_downscaler - -input: - cutout: - - lam_0: - grib: - path: forecaster/20* - pre_processors: - - forward_transform_filter: - # does not have an effect if the temporal downscaler does not use tp as prognostic variable - rescale: - scale: 0.001 # convert from kg m-2 to m - offset: 0 - param: tp - - forward_transform_filter: - rescale: - scale: 0.001 # convert from kg m-2 to m - offset: 0 - param: TOT_PREC - # filters see the model name (tcc) with anemoi-inference <= 0.11.1 and the GRIB name (CLCT) from 0.11.2 - - forward_transform_filter: # cloud cover percent -> fraction, as in training - rescale: - scale: 0.01 - offset: 0 - param: tcc - - forward_transform_filter: # cloud cover percent -> fraction, as in training - rescale: - scale: 0.01 - offset: 0 - param: CLCT - - forward_transform_filter: # cloud cover percent -> fraction, as in training - rescale: - scale: 0.01 - offset: 0 - param: lcc - - forward_transform_filter: # cloud cover percent -> fraction, as in training - rescale: - scale: 0.01 - offset: 0 - param: CLCL - - forward_transform_filter: # cloud cover percent -> fraction, as in training - rescale: - scale: 0.01 - offset: 0 - param: mcc - - forward_transform_filter: # cloud cover percent -> fraction, as in training - rescale: - scale: 0.01 - offset: 0 - param: CLCM - - forward_transform_filter: # cloud cover percent -> fraction, as in training - rescale: - scale: 0.01 - offset: 0 - param: hcc - - forward_transform_filter: # cloud cover percent -> fraction, as in training - rescale: - scale: 0.01 - offset: 0 - param: CLCH - namer: &namer - rules: - # ICON short names (uppercase) - - - shortName: T - - t_{level} - - - shortName: U - - u_{level} - - - shortName: V - - v_{level} - - - shortName: W - - w_{level} - - - shortName: QV - - q_{level} - - - shortName: FI - - z_{level} - - - shortName: PMSL - - msl - - - shortName: FIS - - z - - - shortName: PS - - sp - - - shortName: T_2M - - 2t - - - shortName: TD_2M - - 2d - - - shortName: SKT - - skt - - - shortName: CLCT - - tcc - - - shortName: CLCL - - lcc - - - shortName: CLCM - - mcc - - - shortName: CLCH - - hcc - - - shortName: U_10M - - 10u - - - shortName: V_10M - - 10v - - - shortName: FR_LAND - - lsm - - - shortName: TOT_PREC - - tp - - global: - grib: - path: forecaster/ifs* - namer: *namer - -constant_forcings: - test: - use_original_paths: true - -patch_metadata: resources/sgm-temporal-downscaler-ich1-oper-patch.yaml - -post_processors: - - accumulate_from_start_of_forecast: # accumulate tp from start of forecast - accumulations: - - tp - - forward_transform_filter: - rescale: - scale: 1000 # convert units from m to kg m-2 - offset: 0 - param: tp - - forward_transform_filter: - rescale: - scale: 1000 # convert units from m to kg m-2 - offset: 0 - param: TOT_PREC - -output: - tee: - - fdb: - encoding: - typeOfGeneratingProcess: 2 - generatingProcessIdentifier: 180 - # localNumberOfExperiment: 900 # not supported in current templates; marsClass=od fallback - templates: - samples: resources/templates_index_icon.yaml - post_processors: - - extract_mask: # removes global points - mask: "lam_0/cutout_mask" - as_slice: true - # filters see the model name (tcc) with anemoi-inference <= 0.11.1 and the GRIB name (CLCT) from 0.11.2 - - forward_transform_filter: # cloud cover fraction -> percent, before copy-prognostic - rescale: - scale: 100 - offset: 0 - param: tcc - - forward_transform_filter: # cloud cover fraction -> percent, before copy-prognostic - rescale: - scale: 100 - offset: 0 - param: CLCT - - forward_transform_filter: # cloud cover fraction -> percent, before copy-prognostic - rescale: - scale: 100 - offset: 0 - param: lcc - - forward_transform_filter: # cloud cover fraction -> percent, before copy-prognostic - rescale: - scale: 100 - offset: 0 - param: CLCL - - forward_transform_filter: # cloud cover fraction -> percent, before copy-prognostic - rescale: - scale: 100 - offset: 0 - param: mcc - - forward_transform_filter: # cloud cover fraction -> percent, before copy-prognostic - rescale: - scale: 100 - offset: 0 - param: CLCM - - forward_transform_filter: # cloud cover fraction -> percent, before copy-prognostic - rescale: - scale: 100 - offset: 0 - param: hcc - - forward_transform_filter: # cloud cover fraction -> percent, before copy-prognostic - rescale: - scale: 100 - offset: 0 - param: CLCH - - - fdb: - encoding: - typeOfGeneratingProcess: 2 - generatingProcessIdentifier: 180 - # localNumberOfExperiment is not available in IFS/ECMWF GRIB templates; - # omitting it means marsClass=od (GPI=180 fallback) for the global grid - # Set native marsStream/marsType codetable keys so g2_mars_labeling returns - # stream=enfo and type=fc. Cannot use a concept override in local.98.def - # for marsStream because it conflicts with the codetable key that - # g2_mars_labeling expects as input. Setting them natively is portable - # across eccodes builds. - marsStream: enfo - marsType: fc - templates: - samples: resources/templates_index_ifs.yaml - post_processors: - - extract_mask: # removes lam points - mask: "lam_0/cutout_mask" - as_slice: true - inverse: true - - assign_mask: # fill local/global overlapping points with nan - mask: "global/cutout_mask" - modifiers: - - patches: - - variable: - U_10M: {"param": 165, "shortName": "10u"} - 10u: {"param": 165, "shortName": "10u"} - V_10M: {"param": 166, "shortName": "10v"} - 10v: {"param": 166, "shortName": "10v"} - TD_2M: {"param": 168, "shortName": "2d"} - 2d: {"param": 168, "shortName": "2d"} - T_2M: {"param": 167, "shortName": "2t"} - 2t: {"param": 167, "shortName": "2t"} - FR_LAND: {"param": 172, "shortName": "lsm"} - lsm: {"param": 172, "shortName": "lsm"} - PMSL: {"param": 151, "shortName": "msl"} - msl: {"param": 151, "shortName": "msl"} - PS: {"param": 134, "shortName": "sp"} - sp: {"param": 134, "shortName": "sp"} - SSO_SIGMA: {"param": 163, "shortName": "slor"} - slor: {"param": 163, "shortName": "slor"} - SSO_STDH: {"param": 160, "shortName": "sdor"} - sdor: {"param": 160, "shortName": "sdor"} - CLCT: {"param": 164, "shortName": "tcc"} - tcc: {"param": 164, "shortName": "tcc"} - CLCL: {"param": 186, "shortName": "lcc"} - lcc: {"param": 186, "shortName": "lcc"} - CLCM: {"param": 187, "shortName": "mcc"} - mcc: {"param": 187, "shortName": "mcc"} - CLCH: {"param": 188, "shortName": "hcc"} - hcc: {"param": 188, "shortName": "hcc"} - TOT_PREC: {"param": 228, "shortName": "tp"} - tp: {"param": 228, "shortName": "tp"} - z: {"param": 129, "shortName": "z"} - "^q_(\\d+)$": {"param": 133, "shortName": "q"} - "^t_(\\d+)$": {"param": 130, "shortName": "t"} - "^u_(\\d+)$": {"param": 131, "shortName": "u"} - "^v_(\\d+)$": {"param": 132, "shortName": "v"} - "^w_(\\d+)$": {"param": 135, "shortName": "w"} - "^z_(\\d+)$": {"param": 129, "shortName": "z"} - -# silenced due to bug in anemoi-inference for multi-step temporal downscalers, can be removed when fixed -verbosity: 0 -allow_nans: true -output_frequency: "1h" From e81937ddc6838088150963388b1ad6e02d11b706 Mon Sep 17 00:00:00 2001 From: Carlos Osuna Date: Wed, 21 Oct 2026 14:11:30 +0200 Subject: [PATCH 17/21] merge configs --- config/varda-single-1.0-inference-only.yaml | 2 +- ...tidataset-forecaster-global-ich1-oper.yaml | 2 + ...oral-downscaler-global_trimedge_multi.yaml | 11 ++ src/data_input/__init__.py | 14 ++- src/plotting/compat.py | 117 ++++++++++-------- workflow/rules/inference.smk | 23 +++- workflow/scripts/inference_prepare.py | 49 +++++++- 7 files changed, 162 insertions(+), 56 deletions(-) diff --git a/config/varda-single-1.0-inference-only.yaml b/config/varda-single-1.0-inference-only.yaml index ceabcd56..b465d850 100644 --- a/config/varda-single-1.0-inference-only.yaml +++ b/config/varda-single-1.0-inference-only.yaml @@ -14,7 +14,7 @@ runs: checkpoint: https://service.meteoswiss.ch/mlstore#/models/sruc-m-2-interpolator/versions/3 label: Varda-single-1.0 steps: 0/120/1 - config: resources/inference/configs/sgm-temporal-downscaler-global_trimedge_multi-fdb.yaml + config: resources/inference/configs/sgm-temporal-downscaler-global_trimedge_multi.yaml output: fdb extra_requirements: - anemoi-datasets==0.5.35 diff --git a/resources/inference/configs/sgm-multidataset-forecaster-global-ich1-oper.yaml b/resources/inference/configs/sgm-multidataset-forecaster-global-ich1-oper.yaml index b6a0c4cb..c1546a54 100644 --- a/resources/inference/configs/sgm-multidataset-forecaster-global-ich1-oper.yaml +++ b/resources/inference/configs/sgm-multidataset-forecaster-global-ich1-oper.yaml @@ -32,6 +32,7 @@ output: path: grib/{date}{time:04}_{step:03}.grib encoding: typeOfGeneratingProcess: 2 + generatingProcessIdentifier: 180 centre: lssw templates: samples: resources/templates_index_icon.yaml @@ -84,6 +85,7 @@ output: path: grib/ifs-{date}{time:04}_{step:03}.grib encoding: typeOfGeneratingProcess: 2 + generatingProcessIdentifier: 180 centre: ecmf templates: samples: resources/templates_index_ifs.yaml diff --git a/resources/inference/configs/sgm-temporal-downscaler-global_trimedge_multi.yaml b/resources/inference/configs/sgm-temporal-downscaler-global_trimedge_multi.yaml index cc1be1bb..654f1de4 100644 --- a/resources/inference/configs/sgm-temporal-downscaler-global_trimedge_multi.yaml +++ b/resources/inference/configs/sgm-temporal-downscaler-global_trimedge_multi.yaml @@ -134,6 +134,7 @@ output: path: grib/{dateTime}_{step:03}.grib encoding: typeOfGeneratingProcess: 2 + generatingProcessIdentifier: 180 templates: samples: resources/templates_index_icon.yaml post_processors: @@ -195,6 +196,16 @@ output: path: grib/ifs-{dateTime}_{step:03}.grib encoding: typeOfGeneratingProcess: 2 + generatingProcessIdentifier: 180 + # localNumberOfExperiment is not available in IFS/ECMWF GRIB templates; + # omitting it means marsClass=od (GPI=180 fallback) for the global grid + # Set native marsStream/marsType codetable keys so g2_mars_labeling returns + # stream=enfo and type=fc. Cannot use a concept override in local.98.def + # for marsStream because it conflicts with the codetable key that + # g2_mars_labeling expects as input. Setting them natively is portable + # across eccodes builds. + marsStream: enfo + marsType: fc templates: samples: resources/templates_index_ifs.yaml post_processors: diff --git a/src/data_input/__init__.py b/src/data_input/__init__.py index fbafe71d..dd3434ba 100644 --- a/src/data_input/__init__.py +++ b/src/data_input/__init__.py @@ -124,7 +124,7 @@ def rename_ifs_to_icon(ds: xr.Dataset, mapping: dict[str, str]) -> xr.Dataset: # _FDB_GLOBAL_MODEL ("varda-single-g", IFS-origin global grid) is archived # under standard WMO/ECMWF param ids, not VARDA's custom local table -- see # the `modifiers.patches` block in -# resources/inference/configs/sgm-temporal-downscaler-global_trimedge_multi-fdb.yaml, +# resources/inference/configs/sgm-temporal-downscaler-global_trimedge_multi.yaml, # which explicitly pairs each ICON name with its archived (param, shortName). # No pressure-level params mapped here: nothing in this codebase currently # requests one from the global grid. @@ -595,7 +595,11 @@ def _discover_icon_member_ids( ) -def load_from_grib_file(file: str | list[str], sel_kwargs): +def load_fieldlist_from_grib_file(file: str | list[str], sel_kwargs): + """Load an earthkit FieldList from one or more local GRIB files, filtered + by `sel_kwargs` -- the part of :func:`load_from_grib_file` that's reusable + by callers (e.g. plotting.compat) that need the fieldlist itself rather + than an already-decoded xarray Dataset.""" # Coerce Path objects to str: earthkit-data unwraps a single-element list # into one File source without converting, and then fails on non-str paths. if isinstance(file, (list, tuple)): @@ -603,7 +607,11 @@ def load_from_grib_file(file: str | list[str], sel_kwargs): else: file = str(file) fieldlist = ekd.from_source("file", file, lazily=True).to_fieldlist() - return fieldlist_to_xarray(fieldlist.sel(**sel_kwargs)) + return fieldlist.sel(**sel_kwargs) + + +def load_from_grib_file(file: str | list[str], sel_kwargs): + return fieldlist_to_xarray(load_fieldlist_from_grib_file(file, sel_kwargs)) def variable_name_profile( diff --git a/src/plotting/compat.py b/src/plotting/compat.py index 5ca1e654..82728db3 100644 --- a/src/plotting/compat.py +++ b/src/plotting/compat.py @@ -1,13 +1,12 @@ from datetime import datetime from pathlib import Path -import tempfile - import geopandas as gpd import numpy as np from shapely.geometry import MultiPoint from data_input import ( - load_from_grib_file, + load_fieldlist_from_grib_file, + fieldlist_to_xarray, rename_ifs_to_icon, _fetch_fdb_fieldlist, _FDB_MODEL, @@ -34,28 +33,31 @@ PARAMS_MAP_INV = {v: k for k, v in PARAMS_MAP.items()} -def load_state_from_grib( - file: Path, paramlist: list[str] | None = None +def _decode_state( + regional_fieldlist, + global_fieldlist, + paramlist: list[str] | None, + source: str, ) -> dict[str, np.ndarray | dict[str, np.ndarray] | gpd.GeoSeries]: - # Also request IFS shortname aliases - paramlist_extended = list( - {p for p in (paramlist or [])} - | {PARAMS_MAP[p] for p in (paramlist or []) if p in PARAMS_MAP} - ) - ds = load_from_grib_file(file, {"parameter.variable": paramlist_extended}) + """Build a plotting state dict from the regional (lam_0) and optional + global (IFS) fieldlists for one lead time. + + This is the common decoding step shared by :func:`load_state_from_grib` + and :func:`load_state_from_fdb`, which differ only in how they fetch the + two fieldlists (local files vs FDB retrieve) -- once both have an + earthkit FieldList in hand, the xarray conversion, renaming, and + flat-point-list state assembly are identical. `source` is only used to + make the "nothing found" error message useful. + """ + ds = fieldlist_to_xarray(regional_fieldlist) # Rename any IFS shortnames back to COSMO names ds = rename_ifs_to_icon(ds, PARAMS_MAP_INV) - # TODO check if needed - # if "tp" in ifs_rename and "TOT_PREC" in ds: - # # IFS/ECMWF convention: "tp" is accumulated precip in meters. - # # Convert to kg m-2 (mm) to match the ICON-native convention. - # ds["TOT_PREC"] = ds["TOT_PREC"] * 1000 state = {} ref_param = next((p for p in (paramlist or []) if p in ds), None) if ref_param is None: raise ValueError( - f"None of the requested params {paramlist} found in {file}. " - "The GRIB file may not contain these fields at this lead time " + f"None of the requested params {paramlist} found in {source}. " + "The source may not contain these fields at this lead time " "(e.g. accumulated fields like TOT_PREC are undefined at step 0)." ) # Cast datetime64 directly: datetime.fromtimestamp() would read the epoch @@ -85,24 +87,24 @@ def load_state_from_grib( state["fields"][param] = np.full( ds["longitude"].values.size, np.nan, dtype=float ) - global_file = str(file.parent / f"ifs-{file.stem}.grib") - if Path(global_file).exists(): + + if global_fieldlist is not None: _paramlist_ecmwf = [PARAMS_MAP[p] for p in paramlist] - ds = load_from_grib_file(global_file, {"parameter.variable": _paramlist_ecmwf}) - mask = ~np.isnan(ds[_paramlist_ecmwf[0]].values.squeeze()) - global_lons = ds["longitude"].values.flatten() + gds = fieldlist_to_xarray(global_fieldlist) + mask = ~np.isnan(gds[_paramlist_ecmwf[0]].values.squeeze()) + global_lons = gds["longitude"].values.flatten() if np.max(global_lons) > 180: global_lons = ((global_lons + 180) % 360) - 180 state["longitudes"] = np.concatenate([state["longitudes"], global_lons[mask]]) state["latitudes"] = np.concatenate( - [state["latitudes"], ds["latitude"].values.flatten()[mask]] + [state["latitudes"], gds["latitude"].values.flatten()[mask]] ) for param in _paramlist_ecmwf: - if param in ds: + if param in gds: state["fields"][PARAMS_MAP_INV[param]] = np.concatenate( [ state["fields"][PARAMS_MAP_INV[param]], - ds[param].values.flatten()[mask], + gds[param].values.flatten()[mask], ] ) else: @@ -115,6 +117,29 @@ def load_state_from_grib( return state +def load_state_from_grib( + file: Path, paramlist: list[str] | None = None +) -> dict[str, np.ndarray | dict[str, np.ndarray] | gpd.GeoSeries]: + # Also request IFS shortname aliases + paramlist_extended = list( + {p for p in (paramlist or [])} + | {PARAMS_MAP[p] for p in (paramlist or []) if p in PARAMS_MAP} + ) + regional_fieldlist = load_fieldlist_from_grib_file( + file, {"parameter.variable": paramlist_extended} + ) + + global_fieldlist = None + global_file = file.parent / f"ifs-{file.stem}.grib" + if global_file.exists(): + _paramlist_ecmwf = [PARAMS_MAP[p] for p in paramlist] + global_fieldlist = load_fieldlist_from_grib_file( + global_file, {"parameter.variable": _paramlist_ecmwf} + ) + + return _decode_state(regional_fieldlist, global_fieldlist, paramlist, source=str(file)) + + def load_state_from_fdb( fdb_root: str, reftime: datetime, step: int, paramlist: list[str] | None = None ) -> dict[str, np.ndarray | dict[str, np.ndarray] | gpd.GeoSeries]: @@ -122,39 +147,33 @@ def load_state_from_fdb( Fetches the regional (`_FDB_MODEL`) and global (`_FDB_GLOBAL_MODEL`) grids from FDB -- exactly the requested params, via `_fetch_fdb_fieldlist` -- - and writes them to temp files using the same naming convention - `load_state_from_grib` already expects locally (a primary file plus an - `ifs-.grib` companion in the same directory), then delegates to it - unchanged -- reusing its state-dict construction and global-grid - concatenation (it works with flat point lists, not an aligned xarray grid, - so the regional/global point-count mismatch that :func:`data_input. - load_forecast_data` sidesteps by dropping the global grid entirely isn't a - problem here). + and decodes them directly through the same :func:`_decode_state` shared + with the local-file path (it works with flat point lists, not an aligned + xarray grid, so the regional/global point-count mismatch that + :func:`data_input.load_forecast_data` sidesteps by dropping the global + grid entirely isn't a problem here). """ - regional = _fetch_fdb_fieldlist(fdb_root, reftime, [step], paramlist, model=_FDB_MODEL) - if regional is None: + regional_fieldlist = _fetch_fdb_fieldlist( + fdb_root, reftime, [step], paramlist, model=_FDB_MODEL + ) + if regional_fieldlist is None: raise FileNotFoundError( f"No FDB data found for step={step} model={_FDB_MODEL} in {fdb_root}" ) try: - global_ = _fetch_fdb_fieldlist(fdb_root, reftime, [step], paramlist, model=_FDB_GLOBAL_MODEL) + global_fieldlist = _fetch_fdb_fieldlist( + fdb_root, reftime, [step], paramlist, model=_FDB_GLOBAL_MODEL + ) except ValueError: # Some requested params may have no known id for the global grid # (_FDB_GLOBAL_PARAM_IDS only covers what's actually been needed from # it) -- tolerate that the same way as "nothing archived" below, # rather than failing the whole frame over a companion grid. - global_ = None - - with tempfile.TemporaryDirectory() as tmp_dir: - primary = Path(tmp_dir) / "frame.grib" - with open(primary, "wb") as f: - for fld in regional: - f.write(fld.message()) - if global_ is not None: - with open(Path(tmp_dir) / "ifs-frame.grib", "wb") as f: - for fld in global_: - f.write(fld.message()) - return load_state_from_grib(primary, paramlist=paramlist) + global_fieldlist = None + + return _decode_state( + regional_fieldlist, global_fieldlist, paramlist, source=f"{fdb_root} (step={step})" + ) def load_state_from_raw( diff --git a/workflow/rules/inference.smk b/workflow/rules/inference.smk index 396b1d4b..b7c3aeb3 100644 --- a/workflow/rules/inference.smk +++ b/workflow/rules/inference.smk @@ -211,17 +211,35 @@ def _get_fdb_uenv_for_env(env_id): def _count_fdb_blocks(run_id): - """Count distinct 'fdb:' targets in this run's own inference config output.tee. + """Count how many output.tee targets this run will write into FDB. + + Zero whenever this run doesn't declare output: fdb -- regardless of what + its config's output.tee contains, a grib-output run never writes to FDB. + + For output: fdb runs, inference configs are maintained once, in grib: form, + as the single source of truth for both backends -- inference_prepare.py's + _convert_tee_to_fdb rewrites every 'grib:' tee entry to 'fdb:' at + config-prep time (see resources/inference/configs/ + sgm-temporal-downscaler-global_trimedge_multi.yaml). So each on-disk + 'grib:' entry becomes one FDB target; a literal 'fdb:' entry (a config + authored directly in fdb: form, bypassing the conversion) is also counted, + for robustness. A single init_time can write several output blocks (e.g. an ICON-grid target and an IFS-grid target) into the same env_id-keyed root; the check needs to know how many to expect so a run that silently dropped one entirely doesn't read as complete just because the other is. """ + if RUN_CONFIGS[run_id]["output"] != "fdb": + return 0 with open(RUN_CONFIGS[run_id]["config"]) as f: run_cfg = yaml.safe_load(f) tee = ((run_cfg.get("output") or {}).get("tee")) or [] - return sum(1 for entry in tee if isinstance(entry, dict) and "fdb" in entry) + return sum( + 1 + for entry in tee + if isinstance(entry, dict) and ({"fdb", "grib"} & entry.keys()) + ) checkpoint inference_check_fdb: @@ -416,6 +434,7 @@ rule inference_prepare_temporal_downscaler: if RUN_CONFIGS[wc.run_id].get("forecaster") is None else _get_forecaster_run_id(wc.run_id) ), + output_backend=lambda wc: RUN_CONFIGS[wc.run_id]["output"], script: "../scripts/inference_prepare.py" diff --git a/workflow/scripts/inference_prepare.py b/workflow/scripts/inference_prepare.py index 5b0b075d..462647f0 100644 --- a/workflow/scripts/inference_prepare.py +++ b/workflow/scripts/inference_prepare.py @@ -12,6 +12,7 @@ def prepare_config( default_config_path: str, output_config_path: str, params: dict, + output_backend: str = "grib", ): """Prepare the configuration file for the inference run. @@ -26,11 +27,19 @@ def prepare_config( Path where the updated configuration file will be written. params : dict Dictionary of parameters to override in the default configuration. + output_backend : str + "grib" (default) leaves output.tee untouched -- the inference config is + itself the single source of truth, always written in grib: form. "fdb" + rewrites every `output.tee` entry from `grib:` to `fdb:` (dropping the + grib-only `path` key) -- see `_convert_tee_to_fdb`. """ with open(default_config_path, "r") as f: config = yaml.safe_load(f) + if output_backend == "fdb": + config = _convert_tee_to_fdb(config) + config = _override_recursive(config, params) with open(output_config_path, "w") as f: @@ -111,7 +120,12 @@ def prepare_temporal_downscaler(smk): # prepare config overrides = _overrides_from_params(smk) - prepare_config(smk.input.config, smk.output.config, overrides) + prepare_config( + smk.input.config, + smk.output.config, + overrides, + output_backend=smk.params.output_backend, + ) LOG.info("Wrote config file at %s", smk.output.config) with open(smk.output.config, "r") as f: @@ -177,6 +191,39 @@ def _setup_logger(smk) -> logging.Logger: return LOG +def _convert_tee_to_fdb(config: dict) -> dict: + """Rewrite every `output.tee` entry from `{grib: {...}}` to `{fdb: {...}}`. + + The inference config is maintained once, in grib: form, as the single + source of truth for both backends -- see + resources/inference/configs/sgm-temporal-downscaler-global_trimedge_multi.yaml, + whose output.tee encoding/templates/post_processors/modifiers are already + FDB-compatible (generatingProcessIdentifier, marsStream/marsType, etc.) and + were kept in parity with the (now-deleted) *-fdb-parity.yaml by hand before + this. Only the grib-specific `path` key doesn't apply to fdb: and is + dropped; everything else under the tee entry carries over unchanged. + """ + tee = config.get("output", {}).get("tee") + if not tee: + raise ValueError( + "output_backend='fdb' but this inference config has no output.tee " + "entries to convert -- check the referenced config file." + ) + new_tee = [] + for entry in tee: + if set(entry.keys()) != {"grib"}: + raise ValueError( + f"output_backend='fdb' expects every output.tee entry to be a " + f"single 'grib:' block, got keys {sorted(entry.keys())} -- " + "this config isn't written in the expected grib: source-of-truth form." + ) + inner = dict(entry["grib"]) + inner.pop("path", None) + new_tee.append({"fdb": inner}) + config["output"]["tee"] = new_tee + return config + + def _override_recursive(original: dict, updates: dict) -> dict: """Recursively override values in the original dictionary with those from the updates dictionary.""" for key, value in updates.items(): From 5d048c777e97399006fb631294e438d3f1dc6f85 Mon Sep 17 00:00:00 2001 From: Carlos Osuna Date: Thu, 22 Oct 2026 09:13:23 +0200 Subject: [PATCH 18/21] temporary disable the copy post-processor --- ...tidataset-forecaster-global-ich1-oper.yaml | 9 ++ ...oral-downscaler-global_trimedge_multi.yaml | 33 +++--- src/data_input/__init__.py | 35 +++++- src/evalml/config.py | 19 +-- workflow/rules/inference.smk | 26 +++- workflow/scripts/inference_prepare.py | 112 ++++++++++++++++-- workflow/tools/config.schema.json | 2 +- 7 files changed, 193 insertions(+), 43 deletions(-) diff --git a/resources/inference/configs/sgm-multidataset-forecaster-global-ich1-oper.yaml b/resources/inference/configs/sgm-multidataset-forecaster-global-ich1-oper.yaml index c1546a54..6fc349c2 100644 --- a/resources/inference/configs/sgm-multidataset-forecaster-global-ich1-oper.yaml +++ b/resources/inference/configs/sgm-multidataset-forecaster-global-ich1-oper.yaml @@ -87,6 +87,15 @@ output: typeOfGeneratingProcess: 2 generatingProcessIdentifier: 180 centre: ecmf + # localNumberOfExperiment is not available in IFS/ECMWF GRIB templates; + # omitting it means marsClass=od (GPI=180 fallback) for the global grid + # Set native marsStream/marsType codetable keys so g2_mars_labeling returns + # stream=enfo and type=fc. Cannot use a concept override in local.98.def + # for marsStream because it conflicts with the codetable key that + # g2_mars_labeling expects as input. Setting them natively is portable + # across eccodes builds. + marsStream: enfo + marsType: fc templates: samples: resources/templates_index_ifs.yaml post_processors: diff --git a/resources/inference/configs/sgm-temporal-downscaler-global_trimedge_multi.yaml b/resources/inference/configs/sgm-temporal-downscaler-global_trimedge_multi.yaml index 654f1de4..dd1003de 100644 --- a/resources/inference/configs/sgm-temporal-downscaler-global_trimedge_multi.yaml +++ b/resources/inference/configs/sgm-temporal-downscaler-global_trimedge_multi.yaml @@ -183,14 +183,18 @@ output: offset: 0 param: CLCH # here, the trimedge mask can be specified when available - # at overlap steps (multiples of the forecaster stride) replace the - # re-predicted prognostics with the forecaster's boundary values - - forward_transform_filter: - copy-prognostic-from-forecaster: - forecaster_path: forecaster/20* - common_leadtime: 6h - params_to_keep: [tp] # NOTE, if more diagnostics are added, update accordingly. - namer: *namer + # DISABLED: copy-prognostic-from-forecaster is incompatible with an + # FDB-backed forecaster (it globs the local forecaster/ dir). Waiting + # on the design decision in ecmwf/anemoi-inference#468 for a native, + # backend-agnostic replacement before re-enabling. Until then, + # overlap steps (multiples of the forecaster stride) carry the + # downscaler's own re-predicted prognostics, not the forecaster's. + # - forward_transform_filter: + # copy-prognostic-from-forecaster: + # forecaster_path: forecaster/20* + # common_leadtime: 6h + # params_to_keep: [tp] # NOTE, if more diagnostics are added, update accordingly. + # namer: *namer - grib: path: grib/ifs-{dateTime}_{step:03}.grib @@ -215,12 +219,13 @@ output: inverse: true - assign_mask: # fill local/global overlapping points with nan mask: "global/cutout_mask" - - forward_transform_filter: - copy-prognostic-from-forecaster: - forecaster_path: forecaster/ifs* - common_leadtime: 6h - params_to_keep: [tp] # NOTE, if more diagnostics are added, update accordingly. - namer: *namer + # DISABLED: see note above the lam_0 tee's copy-prognostic-from-forecaster block. + # - forward_transform_filter: + # copy-prognostic-from-forecaster: + # forecaster_path: forecaster/ifs* + # common_leadtime: 6h + # params_to_keep: [tp] # NOTE, if more diagnostics are added, update accordingly. + # namer: *namer modifiers: - patches: - variable: diff --git a/src/data_input/__init__.py b/src/data_input/__init__.py index dd3434ba..5507e811 100644 --- a/src/data_input/__init__.py +++ b/src/data_input/__init__.py @@ -126,8 +126,12 @@ def rename_ifs_to_icon(ds: xr.Dataset, mapping: dict[str, str]) -> xr.Dataset: # the `modifiers.patches` block in # resources/inference/configs/sgm-temporal-downscaler-global_trimedge_multi.yaml, # which explicitly pairs each ICON name with its archived (param, shortName). -# No pressure-level params mapped here: nothing in this codebase currently -# requests one from the global grid. +# Pressure-level ids (T/U/V/QV/FI/OMEGA) match that same file's regex patches +# (e.g. "^t_(\\d+)$" -> param 130) -- needed when the downscaler's own model +# input reads the global grid's pressure levels directly from FDB (see +# inference_prepare.py's _convert_cutout_to_fdb), not just this module's own +# _fetch_fdb_fieldlist reads (which so far have only ever requested surface +# params from the global grid). _FDB_GLOBAL_PARAM_IDS: dict[str, str] = { "PS": "134", "PMSL": "151", @@ -137,12 +141,39 @@ def rename_ifs_to_icon(ds: xr.Dataset, mapping: dict[str, str]) -> xr.Dataset: "V_10M": "166", "TOT_PREC": "228", "FR_LAND": "172", + "T": "130", + "U": "131", + "V": "132", + "QV": "133", + "FI": "129", + "OMEGA": "135", } _FDB_GLOBAL_PL_PARAMS: set[str] = set() _FDB_GLOBAL_PARAM_TIMESPAN: dict[str, str] = { "TOT_PREC": "fs", } +# MARS identity (class/stream/type/expver) that a GPI=180 archive resolves to for +# each model -- confirmed empirically by decoding a real archived message's `mars` +# namespace under the correct ECCODES_DEFINITION_PATH (eccodes-cosmo-mars's +# local.215.def maps GPI=180 to class=ai for the local/ICON table; local.98.def has +# no such concept for the ECMWF/IFS table, so the global grid stays class=od and +# needs marsStream/marsType set natively in its own output encoding -- see +# resources/inference/configs/sgm-temporal-downscaler-global_trimedge_multi.yaml). +# Used both when archiving (inference.smk's fdb5_config.yaml) and when a +# downstream run reads another run's FDB output directly as its own model input +# (inference_prepare.py's _convert_cutout_to_fdb). +_FDB_IDENTITY: dict[str, dict[str, str]] = { + _FDB_MODEL: {"class": "ai", "stream": "oper", "type": "fc", "expver": "0001", "model": _FDB_MODEL}, + _FDB_GLOBAL_MODEL: { + "class": "od", + "stream": "enfo", + "type": "fc", + "expver": "0001", + "model": _FDB_GLOBAL_MODEL, + }, +} + # Per-model (param_ids, pl_params, param_timespan) registries, keyed by the # `model` value used in _fetch_fdb_fieldlist's requests. _FDB_MODEL_REGISTRY: dict[str, tuple[dict[str, str], set[str], dict[str, str]]] = { diff --git a/src/evalml/config.py b/src/evalml/config.py index 530a264f..ef53927e 100644 --- a/src/evalml/config.py +++ b/src/evalml/config.py @@ -197,24 +197,11 @@ class TemporalDownscalerConfig(RunConfig): forecaster: ForecasterConfig | None = Field( None, - description="Configuration for the forecaster run that this temporal downscaler is based on.", + description="Configuration for the forecaster run that this temporal downscaler is based on. " + "If the forecaster declares output: fdb, the downscaler reads its output directly from FDB " + "(inference_prepare.py's _convert_cutout_to_fdb) instead of the local grib symlink.", ) - @model_validator(mode="after") - def _forecaster_output_must_be_grib(self) -> "TemporalDownscalerConfig": - # inference_prepare.py's prepare_temporal_downscaler() hardcodes a symlink to - # the forecaster's local grib directory to chain its output into the downscaler - # as input -- an FDB-backed forecaster would make that symlink point at an empty - # directory, failing confusingly deep inside anemoi-inference rather than here. - if self.forecaster is not None and self.forecaster.output == "fdb": - raise ValueError( - "temporal_downscaler.forecaster.output: fdb' is not supported -- the " - "downscaler reads its forecaster's prognostic state via a local grib " - "symlink (inference_prepare.py), so the forecaster must use output: grib " - "(the default). Only the temporal_downscaler's own output may be fdb." - ) - return self - class BaselineConfig(BaseModel): """Configuration for a single baseline to include in the verification.""" diff --git a/workflow/rules/inference.smk b/workflow/rules/inference.smk index b7c3aeb3..62a64510 100644 --- a/workflow/rules/inference.smk +++ b/workflow/rules/inference.smk @@ -5,6 +5,7 @@ import os from pathlib import Path from datetime import datetime +from types import SimpleNamespace rule inference_get_checkpoint: @@ -386,6 +387,7 @@ rule inference_prepare_forecaster: reftime_to_iso=lambda wc: datetime.strptime( wc.init_time, "%Y%m%d%H%M" ).strftime("%Y-%m-%dT%H:%M"), + output_backend=lambda wc: RUN_CONFIGS[wc.run_id]["output"], script: "../scripts/inference_prepare.py" @@ -395,6 +397,22 @@ def _get_forecaster_run_id(run_id): return RUN_CONFIGS[run_id]["forecaster"]["run_id"] +def _resolve_forecaster_fdb_root(run_id): + """FDB root to read this run's forecaster output from, or None. + + None whenever there's no forecaster, or the forecaster writes local grib + (the default) -- in which case the existing local-grib-symlink input + chaining applies instead. inference_prepare.py's _convert_cutout_to_fdb + only runs when this is non-None. + """ + if RUN_CONFIGS[run_id].get("forecaster") is None: + return None + fct_run_id = _get_forecaster_run_id(run_id) + if RUN_CONFIGS[fct_run_id]["output"] != "fdb": + return None + return _resolve_fdb_root_for_run(fct_run_id) + + # Prepare temporal downscaler for a specific run ID rule inference_prepare_temporal_downscaler: input: @@ -404,8 +422,11 @@ rule inference_prepare_temporal_downscaler: holdout=lambda wc: holdout_dep(wc) if run_uses_holdout(wc.run_id) else [], forecasts=lambda wc: ( [ - OUT_ROOT - / f"logs/inference_execute/{_get_forecaster_run_id(wc.run_id)}-{wc.init_time}.ok" + _get_inference_okfile( + SimpleNamespace( + run_id=_get_forecaster_run_id(wc.run_id), init_time=wc.init_time + ) + ) ] if RUN_CONFIGS[wc.run_id].get("forecaster") is not None else [] @@ -434,6 +455,7 @@ rule inference_prepare_temporal_downscaler: if RUN_CONFIGS[wc.run_id].get("forecaster") is None else _get_forecaster_run_id(wc.run_id) ), + forecaster_fdb_root=lambda wc: _resolve_forecaster_fdb_root(wc.run_id), output_backend=lambda wc: RUN_CONFIGS[wc.run_id]["output"], script: "../scripts/inference_prepare.py" diff --git a/workflow/scripts/inference_prepare.py b/workflow/scripts/inference_prepare.py index 462647f0..0764b48b 100644 --- a/workflow/scripts/inference_prepare.py +++ b/workflow/scripts/inference_prepare.py @@ -6,6 +6,7 @@ from pathlib import Path from evalml.helpers import setup_logger +from data_input import _FDB_MODEL, _FDB_GLOBAL_MODEL, _FDB_PARAM_IDS, _FDB_GLOBAL_PARAM_IDS, _FDB_IDENTITY, _FDB_SCHEMA def prepare_config( @@ -13,6 +14,7 @@ def prepare_config( output_config_path: str, params: dict, output_backend: str = "grib", + forecaster_fdb_root: str | None = None, ): """Prepare the configuration file for the inference run. @@ -32,6 +34,13 @@ def prepare_config( itself the single source of truth, always written in grib: form. "fdb" rewrites every `output.tee` entry from `grib:` to `fdb:` (dropping the grib-only `path` key) -- see `_convert_tee_to_fdb`. + forecaster_fdb_root : str | None + If this run's forecaster declares `output: fdb`, the FDB root its output + was archived into -- rewrites this downscaler's own input.cutout entries + from `grib:` to `fdb:` so it reads the forecaster's output directly from + FDB instead of a local grib symlink. None (default) leaves input.cutout + untouched (the forecaster writes local grib, read via the symlink + prepare_temporal_downscaler already creates). """ with open(default_config_path, "r") as f: @@ -40,6 +49,9 @@ def prepare_config( if output_backend == "fdb": config = _convert_tee_to_fdb(config) + if forecaster_fdb_root is not None: + config = _convert_cutout_to_fdb(config, forecaster_fdb_root) + config = _override_recursive(config, params) with open(output_config_path, "w") as f: @@ -102,14 +114,27 @@ def prepare_temporal_downscaler(smk): # prepare forecaster directory fct_run_id = smk.params.forecaster_run_id + forecaster_fdb_root = smk.params.forecaster_fdb_root if fct_run_id != "null": - fct_workdir = ( - smk.params.output_root / "runs" / fct_run_id / smk.wildcards.init_time - ) - (workdir / "forecaster").symlink_to(fct_workdir / "grib") - LOG.info( - "Created symlink to forecaster grib directory at %s", workdir / "forecaster" - ) + if forecaster_fdb_root is None: + fct_workdir = ( + smk.params.output_root / "runs" / fct_run_id / smk.wildcards.init_time + ) + (workdir / "forecaster").symlink_to(fct_workdir / "grib") + LOG.info( + "Created symlink to forecaster grib directory at %s", workdir / "forecaster" + ) + else: + # Still an output directory of this rule (see inference.smk's + # inference_prepare_temporal_downscaler), so it must exist, but + # unused -- input.cutout reads the forecaster's output straight + # from FDB instead. + (workdir / "forecaster").mkdir(parents=True, exist_ok=True) + LOG.info( + "Forecaster output is in FDB (%s) -- downscaler input.cutout will be " + "rewritten to read it directly, no local grib symlink needed.", + forecaster_fdb_root, + ) else: (workdir / "forecaster").mkdir(parents=True, exist_ok=True) (workdir / "forecaster/.dataset").touch() @@ -125,6 +150,7 @@ def prepare_temporal_downscaler(smk): smk.output.config, overrides, output_backend=smk.params.output_backend, + forecaster_fdb_root=forecaster_fdb_root, ) LOG.info("Wrote config file at %s", smk.output.config) @@ -156,7 +182,12 @@ def prepare_forecaster(smk): stage_holdout_stations(smk, workdir, LOG) overrides = _overrides_from_params(smk) - prepare_config(smk.input.config, smk.output.config, overrides) + prepare_config( + smk.input.config, + smk.output.config, + overrides, + output_backend=smk.params.output_backend, + ) LOG.info("Wrote config file at %s", smk.output.config) with open(smk.output.config, "r") as f: config_content = f.read() @@ -224,6 +255,71 @@ def _convert_tee_to_fdb(config: dict) -> dict: return config +# Which FDB model/param-id table a downscaler input.cutout entry reads from, +# keyed by its grid name -- "lam_0" is the regional/ICON grid (VARDA's own local +# GPI=180 table), "global" is the IFS/global grid (standard WMO/ECMWF param ids). +# See _FDB_IDENTITY (data_input/__init__.py) for the class/stream/type/expver +# each model resolves to. +_CUTOUT_FDB_PARAM_IDS = { + "lam_0": (_FDB_MODEL, _FDB_PARAM_IDS), + "global": (_FDB_GLOBAL_MODEL, _FDB_GLOBAL_PARAM_IDS), +} + + +def _convert_cutout_to_fdb(config: dict, fdb_root: str) -> dict: + """Rewrite the downscaler's input.cutout entries from `{grib: {...}}` to + `{fdb: {...}}`, so it reads its forecaster's output directly from FDB + instead of a local grib symlink. + + `pre_processors`/`namer` under the grib: block carry over unchanged -- they + operate on the already-decoded field regardless of how the raw bytes were + fetched. The grib-only `path` key is dropped. In its place we add + `fdb_config` (pointing at the forecaster's FDB root), `param_id_map` + (translating the checkpoint's own param names -- already rewritten to ICON + shortnames by this config's `patch_metadata`, see anemoi-inference's + Metadata.patch -- to this archive's real param ids), and the static + class/stream/type/expver/model identity, via anemoi-inference's FDBInput + (registered as `fdb`, inputs/fdb.py in the anemoi-inference source): + these keys all pass straight through as constructor kwargs. + """ + cutout = config.get("input", {}).get("cutout") + if not cutout: + raise ValueError( + "forecaster declares output: fdb but this downscaler's input has no " + "input.cutout entries to convert -- check the referenced inference config." + ) + new_cutout = [] + for entry in cutout: + if len(entry) != 1: + raise ValueError(f"Expected a single-key cutout entry, got {sorted(entry.keys())}") + (grid_name, block), = entry.items() + if grid_name not in _CUTOUT_FDB_PARAM_IDS: + raise ValueError( + f"Unknown cutout grid '{grid_name}' -- add its FDB identity to " + "_CUTOUT_FDB_PARAM_IDS before converting it to an fdb: input." + ) + if set(block.keys()) != {"grib"}: + raise ValueError( + f"Expected cutout entry '{grid_name}' to be a single 'grib:' block, " + f"got keys {sorted(block.keys())} -- this config isn't written in " + "the expected grib: source-of-truth form." + ) + inner = dict(block["grib"]) + inner.pop("path", None) + model, param_id_map = _CUTOUT_FDB_PARAM_IDS[grid_name] + inner["fdb_config"] = { + "type": "local", + "engine": "toc", + "schema": str(_FDB_SCHEMA.resolve()), + "spaces": [{"handler": "Default", "roots": [{"path": fdb_root}]}], + } + inner["param_id_map"] = dict(param_id_map) + inner.update(_FDB_IDENTITY[model]) + new_cutout.append({grid_name: {"fdb": inner}}) + config["input"]["cutout"] = new_cutout + return config + + def _override_recursive(original: dict, updates: dict) -> dict: """Recursively override values in the original dictionary with those from the updates dictionary.""" for key, value in updates.items(): diff --git a/workflow/tools/config.schema.json b/workflow/tools/config.schema.json index 50d73668..2bff26fc 100644 --- a/workflow/tools/config.schema.json +++ b/workflow/tools/config.schema.json @@ -1071,7 +1071,7 @@ } ], "default": null, - "description": "Configuration for the forecaster run that this temporal downscaler is based on." + "description": "Configuration for the forecaster run that this temporal downscaler is based on. If the forecaster declares output: fdb, the downscaler reads its output directly from FDB (inference_prepare.py's _convert_cutout_to_fdb) instead of the local grib symlink." } }, "required": [ From bdce91ae62de1c75ba09f6c62cb1d04245fda30c Mon Sep 17 00:00:00 2001 From: Carlos Osuna Date: Thu, 22 Oct 2026 09:37:50 +0200 Subject: [PATCH 19/21] remove the use of uenv --- config/varda-single-1.0-inference-only.yaml | 3 - src/evalml/config.py | 9 -- workflow/rules/inference.smk | 113 ++++++-------------- workflow/tools/config.schema.json | 26 ----- 4 files changed, 34 insertions(+), 117 deletions(-) diff --git a/config/varda-single-1.0-inference-only.yaml b/config/varda-single-1.0-inference-only.yaml index b465d850..fd311e82 100644 --- a/config/varda-single-1.0-inference-only.yaml +++ b/config/varda-single-1.0-inference-only.yaml @@ -30,9 +30,6 @@ locations: output_root: output/ profile: - fdb: - srun_uenv: "fdb/5.19:v2" - srun_view: "realtime" executor: slurm global_resources: gpus: 16 diff --git a/src/evalml/config.py b/src/evalml/config.py index ef53927e..fafd9cf1 100644 --- a/src/evalml/config.py +++ b/src/evalml/config.py @@ -70,15 +70,6 @@ class InferenceResources(BaseModel): None, description="Optional command to prepend to the srun invocation, e.g. 'uenv run --view=realtime fdb/5.19:v2 --'.", ) - srun_uenv: str | None = Field( - None, - description="FDB uenv image to bundle into the venv squashfs, e.g. 'fdb/5.19:v2'. " - "When set, native FDB libraries are extracted from this uenv during venv creation.", - ) - srun_view: str | None = Field( - None, - description="FDB uenv view name to use when bundling FDB libraries, e.g. 'realtime'.", - ) class RunConfig(BaseModel): diff --git a/workflow/rules/inference.smk b/workflow/rules/inference.smk index 62a64510..7e320fc4 100644 --- a/workflow/rules/inference.smk +++ b/workflow/rules/inference.smk @@ -85,10 +85,10 @@ rule inference_extract_requirements: # Prepare the inference environment for a specific checkpoint. The venv is built # in /dev/shm (RAM) to avoid heavy IO on the parallel filesystem, then squashed # to a .squashfs image tracked directly as the rule output. FDB native libraries -# are NOT bundled into this image: the FDB uenv's own squashfs is mounted -# alongside the venv squashfs at run time instead (see inference_execute / -# inference_check_fdb), since squashfs-mount can mount several images at once. -# See https://docs.cscs.ch/guides/storage/#python-virtual-environments-with-uenv. +# are bundled straight into this image via pip (pyfdb pulls in fdb5lib/ +# eccodeslib/eckitlib/metkitlib, all manylinux wheels with vendored .so's) -- +# no separate FDB uenv squashfs needs mounting at run time, matching how the +# main evalml venv already reads FDB in verification.smk/plot.smk. rule inference_prepare_env: input: metadata=OUT_ROOT / "data/runs/{env_id}/anemoi.json", @@ -99,7 +99,7 @@ rule inference_prepare_env: OUT_ROOT / "logs/inference_prepare_env/{env_id}.log", localrule: True params: - fdb_uenv=lambda wc: _get_fdb_uenv_for_env(wc.env_id)[0], + needs_fdb=lambda wc: _env_needs_fdb(wc.env_id), shell: """ ( @@ -115,10 +115,9 @@ rule inference_prepare_env: echo "[$(date)] Installing requirements from {input.requirements}..." uv pip install -r {input.requirements} - # Only the pure-Python / non-uenv pieces are prepared here. Native FDB - # libraries, its view symlinks and its config/schema files stay in the FDB - # uenv's own squashfs, mounted alongside this venv at run time. - if [ -n "{params.fdb_uenv}" ]; then + # pyfdb's own wheel dependencies (fdb5lib etc.) bundle the native FDB + # stack, so this is just another pip install -- no uenv involved. + if [ "{params.needs_fdb}" = "True" ]; then echo "[$(date)] Installing pyfdb (pure-Python FDB binding)..." uv pip install pyfdb @@ -199,16 +198,16 @@ def get_resource(wc, field: str, default): return getattr(fdb, field, default) or default -def _get_fdb_uenv_for_env(env_id): - """Return (fdb_uenv, fdb_view) from profile.fdb.""" - fdb = (config.get("profile") or {}).get("fdb") or {} - if isinstance(fdb, dict): - uenv = fdb.get("srun_uenv", "") or "" - view = fdb.get("srun_view", "") or "realtime" - else: - uenv = getattr(fdb, "srun_uenv", "") or "" - view = getattr(fdb, "srun_view", "") or "realtime" - return uenv, view +def _env_needs_fdb(env_id): + """Whether any run sharing this env_id declares output: fdb -- i.e. whether + this shared venv needs the FDB native stack installed. pyfdb bundles its own + native libraries (fdb5lib/eccodeslib/eckitlib/metkitlib wheels), so this is + a plain pip install -- no uenv/squashfs mount of a separate native image + needed.""" + return any( + run_cfg["env_id"] == env_id and run_cfg.get("output") == "fdb" + for run_cfg in RUN_CONFIGS.values() + ) def _count_fdb_blocks(run_id): @@ -251,10 +250,9 @@ checkpoint inference_check_fdb: directory). A 'hit' requires every requested lead time to be present in that root -- see inference_check_fdb.py for why no further MARS-key disambiguation (class/expver/model) is needed. Writes 'hit\\t' or - 'miss\\t' to the status file. Runs inside the squashfs venv so that pyfdb and the - native FDB C library are available. Skipped entirely (writes 'miss') unless this - run's own config declares `output: fdb` (see the fail-fast validation below this - rule, which rejects output: fdb runs whose environment has no FDB uenv bundled). + 'miss\\t' to the status file. Runs inside the squashfs venv so that pyfdb (and + the native FDB C library it bundles) is available. Skipped entirely (writes + 'miss') unless this run's own config declares `output: fdb`. """ input: image=lambda wc: OUT_ROOT @@ -268,10 +266,8 @@ checkpoint inference_check_fdb: params: image_path=lambda wc, input: str(Path(input.image).resolve()), fdb_configured=lambda wc: RUN_CONFIGS[wc.run_id]["output"] == "fdb", - fdb_uenv=lambda wc: _get_fdb_uenv_for_env(RUN_CONFIGS[wc.run_id]["env_id"])[0], fdb_root=lambda wc: _resolve_fdb_root_for_run(wc.run_id), fdb_schema=str(Path("resources/fdb/realtime-varda.schema").resolve()), - fdb_view=lambda wc: get_resource(wc, "srun_view", "realtime") or "realtime", date=lambda wc: datetime.strptime(wc.init_time, "%Y%m%d%H%M").strftime("%Y%m%d"), time=lambda wc: datetime.strptime(wc.init_time, "%Y%m%d%H%M").strftime("%H%M"), steps=lambda wc: RUN_CONFIGS[wc.run_id]["steps"], @@ -284,17 +280,12 @@ checkpoint inference_check_fdb: printf 'miss\t' > {output.status} echo "FDB check skipped: FDB not configured for this environment" else - # Mount the venv (pyfdb, python) and the FDB uenv (native libs, config, - # at its required /user-environment path) side by side in one call. - # squashfs-mount needs the mountpoint dir to already exist -- /user-environment - # is a pre-provisioned system directory, but there's no such slot for a second - # image, so the venv gets a job-local directory we create ourselves. + # Mount just the venv -- pyfdb bundles its own native FDB libraries + # (fdb5lib etc.), no separate uenv image needed. VENV_MOUNT=$(mktemp -d) trap "rmdir '$VENV_MOUNT' 2>/dev/null || true" EXIT - FDB_SQUASHFS=$(uenv image inspect --format='{{sqfs}}' {params.fdb_uenv}) - squashfs-mount -s "{params.image_path}:$VENV_MOUNT,$FDB_SQUASHFS:/user-environment" -- bash -c ' + squashfs-mount -s "{params.image_path}:$VENV_MOUNT" -- bash -c ' source '"$VENV_MOUNT"'/bin/activate - export LD_LIBRARY_PATH=/user-environment/env/{params.fdb_view}/lib64:/user-environment/env/{params.fdb_view}/lib${{LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}} python {input.script} \ --fdb-root {params.fdb_root} \ --fdb-schema {params.fdb_schema} \ @@ -332,8 +323,7 @@ def _get_inference_status(wc): local-grib output), in which case downstream code reads local grib as before. """ - env_id = RUN_CONFIGS[wc.run_id]["env_id"] - if _get_fdb_uenv_for_env(env_id)[0]: + if RUN_CONFIGS[wc.run_id]["output"] == "fdb": status_file = checkpoints.inference_check_fdb.get( run_id=wc.run_id, init_time=wc.init_time ).output.status @@ -504,8 +494,6 @@ rule inference_execute: "disable_local_eccodes_definitions", False ), srun_prefix=lambda wc: get_resource(wc, "srun_prefix", "") or "", - fdb_view=lambda wc: get_resource(wc, "srun_view", "realtime") or "realtime", - fdb_uenv=lambda wc: _get_fdb_uenv_for_env(RUN_CONFIGS[wc.run_id]["env_id"])[0], fdb_configured=lambda wc: RUN_CONFIGS[wc.run_id]["output"] == "fdb", fdb_root=lambda wc: _resolve_fdb_root_for_run(wc.run_id), fdb_schema=str(Path("resources/fdb/realtime-varda.schema").resolve()), @@ -545,24 +533,16 @@ FDBEOF _run_inference() {{ local VENV=$1 - local FDB_MOUNT=${{2:-}} source "$VENV/bin/activate" if [ "{params.disable_local_definitions}" = "False" ]; then export ECCODES_DEFINITION_PATH="$VENV/share/eccodes-cosmo-resources/definitions" fi - # Set up FDB env vars when the FDB uenv is mounted alongside this venv - # (as a second image at $FDB_MOUNT -- see the squashfs-mount call below). - if [ -n "$FDB_MOUNT" ]; then - export LD_LIBRARY_PATH="$FDB_MOUNT/env/{params.fdb_view}/lib64:$FDB_MOUNT/env/{params.fdb_view}/lib${{LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}}" - # FDB5_CONFIG_OVERRIDE is set whenever FDB is configured for this env - # (see fdb_configured above); the uenv's own config is only a fallback. - if [ -n "$FDB5_CONFIG_OVERRIDE" ]; then - export FDB5_CONFIG_FILE="$FDB5_CONFIG_OVERRIDE" - else - export FDB5_CONFIG_FILE="$FDB_MOUNT/meta/recipe/meta/private/fdb_config/{params.fdb_view}.yaml" - fi + # pyfdb bundles its own native FDB libraries -- no uenv mount, no + # LD_LIBRARY_PATH needed. Just point it at this run's FDB config/schema. + if [ "{params.fdb_configured}" = "True" ]; then + export FDB5_CONFIG_FILE="$FDB5_CONFIG_OVERRIDE" # Prepend eccodes-cosmo-mars definitions so MARS namespace concepts # (marsModel, marsStream, marsClass, etc.) resolve correctly for MeteoSwiss GRIBs. # Must be set in ECCODES_DEFINITION_PATH (not just GRIB_DEFINITION_PATH) because @@ -620,27 +600,11 @@ FDBEOF # squashfs-mount needs the mountpoint dir to already exist, and -- since srun # may dispatch to a different physical node -- it must be on the shared - # filesystem, not /tmp (node-local). /user-environment is a pre-provisioned - # system directory present on every node for the FDB uenv's own absolute - # internal symlinks; the venv gets a job-local dir here instead. - # - # srun must be the OUTER call, dispatching a (dual) squashfs-mount as its - # task, not the other way around: wrapping srun inside a dual-image - # squashfs-mount breaks Slurm's plugin initialization on this system - # ("Plugin initialization failed") -- verified empirically. A single-image - # mount around srun is fine; two images at once are not. + # filesystem, not /tmp (node-local). VENV_MOUNT="{params.workdir}/.venv-mount" mkdir -p "$VENV_MOUNT" trap "rmdir '$VENV_MOUNT' 2>/dev/null || true" EXIT - MOUNT_ARGS=("{params.env_path}:$VENV_MOUNT") - INNER_CMD="_run_inference $VENV_MOUNT" - if [ "{params.fdb_configured}" = "True" ]; then - FDB_SQUASHFS=$(uenv image inspect --format='{{sqfs}}' {params.fdb_uenv}) - MOUNT_ARGS+=("$FDB_SQUASHFS:/user-environment") - INNER_CMD="$INNER_CMD /user-environment" - fi - {params.srun_prefix} srun \ --job-name=anemoi-inference \ --unbuffered \ @@ -651,19 +615,16 @@ FDBEOF --time={resources.runtime} \ --gres={resources.gres} \ --ntasks={resources.ntasks} \ - squashfs-mount -s "$(IFS=,; echo "${{MOUNT_ARGS[*]}}")" -- bash -c "$INNER_CMD" + squashfs-mount -s "{params.env_path}:$VENV_MOUNT" -- bash -c "_run_inference $VENV_MOUNT" ) >{log} 2>&1 && touch {output.okfile} """ # fmt: on -# Fail fast: a run declaring `output: fdb` is relying on infrastructure and on its -# own referenced inference config to actually deliver FDB output -- catch the two -# ways that silently doesn't happen, at workflow-build time, before any (expensive) +# Fail fast: a run declaring `output: fdb` is relying on its own referenced +# inference config to actually deliver FDB output -- catch the one way that +# silently doesn't happen, at workflow-build time, before any (expensive) # inference job runs: -# - no FDB uenv configured -> inference_check_fdb/inference_execute's own -# fdb_configured gate would just read False and FDB would never be touched at all, -# despite the declared intent. # - the referenced inference config has zero `fdb:` output.tee blocks -> # inference_check_fdb.py's completeness check is vacuously true when # expected_blocks == 0 (0 complete blocks found >= 0 expected), reporting a false @@ -672,12 +633,6 @@ FDBEOF for _run_id, _run_cfg in RUN_CONFIGS.items(): if _run_cfg.get("output") != "fdb": continue - if not _get_fdb_uenv_for_env(_run_cfg["env_id"])[0]: - raise ValueError( - f"Run '{_run_id}' declares output: fdb but profile.fdb.srun_uenv is not " - "set -- FDB archival/reads need a mounted FDB uenv to supply the native " - "libfdb5 library. Set profile.fdb.srun_uenv (e.g. 'fdb/5.21:v2')." - ) if _count_fdb_blocks(_run_id) == 0: raise ValueError( f"Run '{_run_id}' declares output: fdb but its inference config " diff --git a/workflow/tools/config.schema.json b/workflow/tools/config.schema.json index 2bff26fc..6e446cf1 100644 --- a/workflow/tools/config.schema.json +++ b/workflow/tools/config.schema.json @@ -649,32 +649,6 @@ "default": null, "description": "Optional command to prepend to the srun invocation, e.g. 'uenv run --view=realtime fdb/5.19:v2 --'.", "title": "Srun Prefix" - }, - "srun_uenv": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "default": null, - "description": "FDB uenv image to bundle into the venv squashfs, e.g. 'fdb/5.19:v2'. When set, native FDB libraries are extracted from this uenv during venv creation.", - "title": "Srun Uenv" - }, - "srun_view": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "default": null, - "description": "FDB uenv view name to use when bundling FDB libraries, e.g. 'realtime'.", - "title": "Srun View" } }, "title": "InferenceResources", From df60d04c084a4b5b0b18e20d776c615fee9144e0 Mon Sep 17 00:00:00 2001 From: Carlos Osuna Date: Mon, 2 Nov 2026 11:37:01 +0100 Subject: [PATCH 20/21] check for duplicates --- workflow/rules/inference.smk | 58 +++++++++++++++++++++++++++++++++++- 1 file changed, 57 insertions(+), 1 deletion(-) diff --git a/workflow/rules/inference.smk b/workflow/rules/inference.smk index 7e320fc4..2aeca846 100644 --- a/workflow/rules/inference.smk +++ b/workflow/rules/inference.smk @@ -322,6 +322,13 @@ def _get_inference_status(wc): It's None only when this run's config has no FDB targets at all (plain local-grib output), in which case downstream code reads local grib as before. + + For output: fdb runs that actually execute (a 'miss'), the returned okfile + is inference_check_fdb_duplicates's own -- not inference_execute's -- + so nothing downstream (verification, plotting, a chained downscaler's + forecaster input) consumes this run's FDB data until it's confirmed free + of MARS-key collisions. A cache 'hit' skips this: that data was already + checked when it was first produced. """ if RUN_CONFIGS[wc.run_id]["output"] == "fdb": status_file = checkpoints.inference_check_fdb.get( @@ -330,7 +337,11 @@ def _get_inference_status(wc): status, _, root = Path(status_file).read_text().strip().partition("\t") if status == "hit": return str(status_file), root - okfile = str(OUT_ROOT / f"logs/inference_execute/{wc.run_id}-{wc.init_time}.ok") + okfile = str( + OUT_ROOT / f"logs/inference_check_fdb_duplicates/{wc.run_id}-{wc.init_time}.ok" + ) + else: + okfile = str(OUT_ROOT / f"logs/inference_execute/{wc.run_id}-{wc.init_time}.ok") fdb_root = ( _resolve_fdb_root_for_run(wc.run_id) if _count_fdb_blocks(wc.run_id) > 0 else None ) @@ -621,6 +632,51 @@ FDBEOF # fmt: on +rule inference_check_fdb_duplicates: + """Fail if this run's own FDB archival collided with an already-present + record -- i.e. a MARS key with more than one entry (live + masked) in + FDB's toc index for this init time (see inference_check_fdb_duplicates.py + for why this can only be caught by inspecting the archive after the fact, + not at write time). Only runs for output: fdb runs, right after + inference_execute actually writes -- not on a cache hit (status == "hit" + in _get_inference_status), whose data was already checked when it was + first produced. + """ + input: + image=lambda wc: OUT_ROOT + / f"data/runs/{RUN_CONFIGS[wc.run_id]['env_id']}/venv.squashfs", + okfile=OUT_ROOT / "logs/inference_execute/{run_id}-{init_time}.ok", + script="workflow/scripts/inference_check_fdb_duplicates.py", + output: + okfile=OUT_ROOT / "logs/inference_check_fdb_duplicates/{run_id}-{init_time}.ok", + log: + OUT_ROOT / "logs/inference_check_fdb_duplicates/{run_id}-{init_time}.log", + localrule: True + params: + image_path=lambda wc, input: str(Path(input.image).resolve()), + fdb_root=lambda wc: _resolve_fdb_root_for_run(wc.run_id), + fdb_schema=str(Path("resources/fdb/realtime-varda.schema").resolve()), + date=lambda wc: datetime.strptime(wc.init_time, "%Y%m%d%H%M").strftime("%Y%m%d"), + time=lambda wc: datetime.strptime(wc.init_time, "%Y%m%d%H%M").strftime("%H%M"), + shell: + """ + ( + set -euo pipefail + VENV_MOUNT=$(mktemp -d) + trap "rmdir '$VENV_MOUNT' 2>/dev/null || true" EXIT + squashfs-mount -s "{params.image_path}:$VENV_MOUNT" -- bash -c ' + source '"$VENV_MOUNT"'/bin/activate + python {input.script} \ + --fdb-root {params.fdb_root} \ + --fdb-schema {params.fdb_schema} \ + --date {params.date} \ + --time {params.time} \ + --output {output.okfile} + ' + ) >{log} 2>&1 + """ + + # Fail fast: a run declaring `output: fdb` is relying on its own referenced # inference config to actually deliver FDB output -- catch the one way that # silently doesn't happen, at workflow-build time, before any (expensive) From bbfeb3a8482290f5d6efacdf4a071e000f778f14 Mon Sep 17 00:00:00 2001 From: Carlos Osuna Date: Thu, 12 Nov 2026 08:31:58 +0100 Subject: [PATCH 21/21] remove line from gitignore --- .gitignore | 6 ------ 1 file changed, 6 deletions(-) diff --git a/.gitignore b/.gitignore index e15e0515..07a194b4 100644 --- a/.gitignore +++ b/.gitignore @@ -58,9 +58,3 @@ uv.lock # Core dumps core_* - -# Local clone of MeteoSwiss's eccodes-cosmo-mars (MARS-namespace GRIB -# definitions needed for FDB reads outside the squashfs-mounted venv, see -# workflow/rules/plot.smk and verification.smk). Not vendored: each checkout -# clones it locally, same as the one baked into the inference venv. -/eccodes-cosmo-mars/