diff --git a/README.md b/README.md index 59de1dc..3a66aab 100644 --- a/README.md +++ b/README.md @@ -245,6 +245,8 @@ See the `example` recipe for a complete working example. Key steps: ``` 4. **Run it** with `nflow run` or `nflow run-all` +**Dependencies between stages:** When a stage lists yours in `dependencies`, the runner passes it your stage's experiment names as `run_after`. By default these are `expname`, or `{expname}-{env}` for each selected environment when the stage config defines `environments`. If `execute()` submits experiments under other names, override the `submitted_expnames(config, expname)` classmethod to return the ones dependent stages should wait for (see `nvflow/recipes/finance/stages/rl/validate_questions.py`). + **Terminal output in stages:** Use the `console` helpers for consistent, readable logs when your stage runs (e.g. `console.status()`, `console.detail()`, `console.success()`). See **[Console UI guide](docs/development/console-ui.md)**. Example: `nvflow/recipes/example/stages/sdg/generate_answer.py` diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index dfc3fa3..7cab4b6 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -140,6 +140,7 @@ classDiagram +execute(config, cluster, expname, run_after)* +validate_config(config) +get_dependencies(config) + +submitted_expnames(config, expname)$ } class StageRegistry { @@ -392,7 +393,7 @@ graph LR **Dependency Handling:** 1. WorkflowRunner reads `dependencies` from stage config -2. Generates experiment names for all dependencies +2. Asks each dependency stage for the experiment names it submits (`BaseStage.submitted_expnames`) 3. Passes as `run_after` parameter to dependent stage 4. Slurm uses job dependencies to ensure correct execution order diff --git a/nvflow/core/base_stage.py b/nvflow/core/base_stage.py index c19ae56..b663b18 100644 --- a/nvflow/core/base_stage.py +++ b/nvflow/core/base_stage.py @@ -98,3 +98,40 @@ def get_dependencies(self, config: dict[str, Any]) -> list[str]: List of stage names this stage depends on """ return config.get("dependencies", []) + + @classmethod + def submitted_expnames(cls, config: dict[str, Any], expname: str) -> list[str]: + """Get the experiment names that dependent stages must wait for. + + WorkflowRunner calls this on the stage class, without instantiating it, + and passes the names as ``run_after`` to every stage that lists this + stage in ``dependencies``. Each returned name must identify an + experiment that :meth:`execute` submits for the same ``config`` and + ``expname``. If :meth:`execute` chains experiments, return the + terminal experiment of each chain. + + The default follows the convention most stages use: one experiment + named ``expname``, or, when ``config`` defines ``environments``, one + experiment per selected environment named ``f"{expname}-{env}"``. + Override it, as a classmethod, when :meth:`execute` names its + experiments differently. + + Args: + config: Stage configuration as passed to :meth:`execute`, including + the ``_environment`` filter added by the runner + expname: Experiment name passed to :meth:`execute` + + Returns: + Experiment names for dependent stages to use as ``run_after`` + """ + environments = config.get("environments") + if not environments: + return [expname] + selected = config.get("_environment") + if not selected: + env_names = list(environments) + else: + if isinstance(selected, str): + selected = [selected] + env_names = [env for env in selected if env in environments] + return [f"{expname}-{env}" for env in env_names] diff --git a/nvflow/core/workflow_runner.py b/nvflow/core/workflow_runner.py index c425f50..b67dcf7 100644 --- a/nvflow/core/workflow_runner.py +++ b/nvflow/core/workflow_runner.py @@ -323,10 +323,7 @@ def _run_stage( """ section(f"Running Stage: {stage_name}") - # Get stage configuration and inject environment filter - stage_config = {**self.config["stages"][stage_name]} - if environment is not None: - stage_config["_environment"] = environment + stage_config = self._stage_config(stage_name, environment) # Get stage class from hierarchical registry with explicit context if not StageRegistry.has(self.recipe, self.workflow_name, stage_name): @@ -363,6 +360,13 @@ def _run_stage( success(f"Stage '{stage_name}' completed") + def _stage_config(self, stage_name: str, environment: list[str] | None) -> dict: + """Return the config passed to a stage, including the environment filter.""" + stage_config = {**self.config["stages"][stage_name]} + if environment is not None: + stage_config["_environment"] = environment + return stage_config + def _get_expname(self, stage_name: str, stage_config: dict) -> str: """Generate clean experiment name for a stage. @@ -395,45 +399,21 @@ def _get_run_after_names( ) -> list[str] | None: """Build ``run_after`` experiment names for Slurm dependency tracking. - Per-environment stages submit jobs with ``{expname}-{env_name}`` - suffixes. This method expands dependency names to match those - suffixed experiment names so that ``nemo-run`` can resolve the - correct Slurm job handles. - - For stages without ``environments``, the base experiment name is - used (unchanged from previous behaviour). + Each dependency stage reports the experiments it submits through + :meth:`BaseStage.submitted_expnames`, given the same config and + experiment name that :meth:`_run_stage` passes to its ``execute()``, + so that ``nemo-run`` can resolve the correct Slurm job handles. """ if not dependencies: return None names: list[str] = [] for dep in dependencies: - dep_config = self.config["stages"][dep] - base = self._get_expname(dep, dep_config) - if dep_config.get("environments"): - env_names = self._resolve_env_names(dep_config, environment) - names.extend(f"{base}-{env}" for env in env_names) - else: - names.append(base) + dep_config = self._stage_config(dep, environment) + dep_expname = self._get_expname(dep, dep_config) + dep_class = StageRegistry.get(self.recipe, self.workflow_name, dep) + names.extend(dep_class.submitted_expnames(dep_config, dep_expname)) return names or None - @staticmethod - def _resolve_env_names( - stage_config: dict, - environment: list[str] | None, - ) -> list[str]: - """Return the environment names a stage will iterate over. - - Mirrors the filtering logic of ``resolve_environments()`` in - ``nvflow.lib.rl.helpers`` but operates on the raw config dict - so the core module stays independent of recipe-specific code. - """ - envs = stage_config.get("environments", {}) - if not envs: - return [] - if environment: - return [e for e in environment if e in envs] - return list(envs.keys()) - def _validate_stages(self, stages_to_run: list[str], all_stages: list[str]) -> None: """Validate that requested stages exist and are registered. diff --git a/nvflow/recipes/finance/stages/rl/training.py b/nvflow/recipes/finance/stages/rl/training.py index 64b0a37..1dde89d 100644 --- a/nvflow/recipes/finance/stages/rl/training.py +++ b/nvflow/recipes/finance/stages/rl/training.py @@ -1018,6 +1018,17 @@ def execute( self._display_grpo_summary(prepared, combined_config) self._submit_grpo_job(prepared, cluster_config, combined_config) + @classmethod + def submitted_expnames(cls, config: dict[str, Any], expname: str) -> list[str]: + """Get the single training experiment submitted by :meth:`execute`. + + One environment trains as ``{expname}-{env}``; several environments + train together as ``{expname}-{env1}+{env2}+...``. + """ + from nvflow.lib.rl.helpers import resolve_environments + + return [f"{expname}-{'+'.join(resolve_environments(config))}"] + def validate_config(self, config: dict[str, Any]) -> None: """Validate configuration.""" required = ["output_dir", "model_name", "data_source_dir"] diff --git a/nvflow/recipes/finance/stages/rl/validate_questions.py b/nvflow/recipes/finance/stages/rl/validate_questions.py index 92b5aad..768e2e2 100644 --- a/nvflow/recipes/finance/stages/rl/validate_questions.py +++ b/nvflow/recipes/finance/stages/rl/validate_questions.py @@ -168,6 +168,13 @@ def execute( console.success(f"validate_questions submitted for '{env_name}' -> {final_file}") + @classmethod + def submitted_expnames(cls, config: dict[str, Any], expname: str) -> list[str]: + """Get each environment's phase-2 experiment, which runs after phase 1.""" + from nvflow.lib.rl.helpers import resolve_environments + + return [f"{expname}-{env_name}-phase2-llm" for env_name in resolve_environments(config)] + def validate_config(self, config: dict[str, Any]) -> None: """Basic sanity checks before Slurm submission.""" for field in ("output_dir", "source_data", "prompt_config", "environments"): diff --git a/tests/test_run_after_expnames.py b/tests/test_run_after_expnames.py new file mode 100644 index 0000000..9426ab7 --- /dev/null +++ b/tests/test_run_after_expnames.py @@ -0,0 +1,232 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Tests that ``run_after`` names match the experiments dependency stages submit. + +When a ``run_after`` name does not identify a submitted experiment, nemo-skills +logs a warning and submits the dependent job without that dependency. +""" + +from __future__ import annotations + +import sys +import types +from pathlib import Path +from typing import Any + +import pytest +import yaml + +import nvflow.recipes # noqa: F401 # registers recipe stages +from nvflow.core import BaseStage, StageRegistry +from nvflow.core.workflow_runner import WorkflowRunner + +GRPO_WORKFLOW = ( + Path(__file__).resolve().parent.parent + / "nvflow/recipes/finance/workflows/grpo/qwen3_30b_a3b.yaml" +) + +# Snapshot at collection time: tests/test_core.py calls StageRegistry.clear(). +_DISCOVERED_STAGES = { + recipe: {workflow: dict(stages) for workflow, stages in workflows.items()} + for recipe, workflows in StageRegistry._stages.items() +} + + +@pytest.fixture +def registry(monkeypatch): + """Restore the stages registered by recipe discovery for one test.""" + stages = { + recipe: {workflow: dict(entries) for workflow, entries in workflows.items()} + for recipe, workflows in _DISCOVERED_STAGES.items() + } + monkeypatch.setattr(StageRegistry, "_stages", stages) + return StageRegistry + + +def _stub_modules(monkeypatch, modules: dict[str, dict[str, Any]]) -> None: + """Install fake modules, and their parent packages, for one test.""" + created: dict[str, types.ModuleType] = {} + for dotted, attributes in modules.items(): + parts = dotted.split(".") + for depth in range(1, len(parts) + 1): + name = ".".join(parts[:depth]) + if name not in created: + created[name] = types.ModuleType(name) + monkeypatch.setitem(sys.modules, name, created[name]) + for attribute, value in attributes.items(): + setattr(created[dotted], attribute, value) + + +class _DefaultNamingStage(BaseStage): + def execute(self, config, cluster, expname, run_after=None): + pass + + +@pytest.mark.parametrize( + ("config", "expected"), + [ + ({}, ["wf-stage"]), + ({"environments": {}}, ["wf-stage"]), + ({"environments": {"a": {}, "b": {}}}, ["wf-stage-a", "wf-stage-b"]), + ({"environments": {"a": {}, "b": {}}, "_environment": []}, ["wf-stage-a", "wf-stage-b"]), + ({"environments": {"a": {}, "b": {}}, "_environment": ["b"]}, ["wf-stage-b"]), + ( + {"environments": {"a": {}, "b": {}}, "_environment": ["b", "a"]}, + ["wf-stage-b", "wf-stage-a"], + ), + ({"environments": {"a": {}}, "_environment": ["a", "other"]}, ["wf-stage-a"]), + ({"environments": {"a": {}, "b": {}}, "_environment": "b"}, ["wf-stage-b"]), + ], +) +def test_default_submitted_expnames(config, expected): + assert _DefaultNamingStage.submitted_expnames(config, "wf-stage") == expected + + +def test_run_after_uses_names_reported_by_the_dependency(tmp_path, monkeypatch): + import nvflow.lib.sbatch as sbatch + + calls: dict[str, Any] = {"producer_instances": 0} + + class Producer(BaseStage): + def __init__(self): + calls["producer_instances"] += 1 + + def execute(self, config, cluster, expname, run_after=None): + calls["execute"] = (config, expname) + + @classmethod + def submitted_expnames(cls, config, expname): + calls["submitted_expnames"] = (config, expname) + return [f"{expname}-last"] + + class Consumer(BaseStage): + def execute(self, config, cluster, expname, run_after=None): + calls["run_after"] = run_after + + # run() patches nemo-skills' executor factory, which is not under test here. + monkeypatch.setattr(sbatch, "apply_sbatch_args_autopatch", lambda: None) + monkeypatch.setattr( + StageRegistry, "_stages", {"test": {"wf": {"producer": Producer, "consumer": Consumer}}} + ) + workflow = tmp_path / "workflow.yaml" + workflow.write_text( + yaml.safe_dump( + { + "recipe": "test", + "workflow": {"name": "wf"}, + "cluster": "local", + "environments": {"a": {}, "b": {}}, + "pipeline_stages": ["producer", "consumer"], + "stages": { + "producer": {"environments": {"a": {}, "b": {}}}, + "consumer": {"dependencies": ["producer"]}, + }, + } + ) + ) + + WorkflowRunner(str(workflow)).run(environment=["b"]) + + assert calls["run_after"] == ["wf-producer-last"] + assert calls["submitted_expnames"] == calls["execute"] + # Only _run_stage instantiates the stage; dependency naming is class-level. + assert calls["producer_instances"] == 1 + + +@pytest.mark.parametrize( + "environment", + [None, ["equivalence_llm_judge"], ["equivalence_llm_judge", "mcqa"]], + ids=["all-environments", "one-environment", "two-environments"], +) +def test_grpo_eval_waits_on_the_training_experiment(registry, monkeypatch, environment): + """Training submits one experiment for all selected environments.""" + training_cls = registry.get("finance", "grpo", "training") + eval_cls = registry.get("finance", "grpo", "eval") + submitted: list[str] = [] + received: dict[str, Any] = {} + + _stub_modules( + monkeypatch, + {"nemo_skills.pipeline.utils.cluster": {"get_cluster_config": lambda cluster: {}}}, + ) + # Record the experiment name that _submit_grpo_job passes to get_exp(). + monkeypatch.setattr( + training_cls, + "_prepare_grpo_config", + lambda self, config, cluster, expname, **kwargs: types.SimpleNamespace(expname=expname), + ) + monkeypatch.setattr(training_cls, "_display_grpo_summary", lambda self, *args: None) + monkeypatch.setattr( + training_cls, + "_submit_grpo_job", + lambda self, prepared, *args: submitted.append(prepared.expname), + ) + monkeypatch.setattr(eval_cls, "validate_config", lambda self, config: None) + monkeypatch.setattr( + eval_cls, + "execute", + lambda self, config, cluster, expname, run_after=None: received.update(run_after=run_after), + ) + + runner = WorkflowRunner(str(GRPO_WORKFLOW)) + stages = ["training", "eval"] + runner._run_stage("training", environment=environment, stages_to_run=stages) + runner._run_stage("eval", environment=environment, stages_to_run=stages) + + assert len(submitted) == 1 + assert received["run_after"] == submitted + + +@pytest.mark.parametrize( + "environment", + [None, ["equivalence_llm_judge"]], + ids=["all-environments", "one-environment"], +) +def test_grpo_data_transformation_waits_on_validate_questions_final_jobs( + registry, monkeypatch, environment +): + """validate_questions chains phase 1 -> phase 2 per environment.""" + transform_cls = registry.get("finance", "grpo", "data_transformation") + submissions: list[tuple[str, list[str]]] = [] + received: dict[str, Any] = {} + + def record(**kwargs): + submissions.append((kwargs["expname"], list(kwargs.get("run_after") or []))) + + _stub_modules( + monkeypatch, + { + "nemo_skills.pipeline.cli": { + "generate": record, + "run_cmd": record, + "wrap_arguments": lambda text: text, + } + }, + ) + monkeypatch.setattr(transform_cls, "validate_config", lambda self, config: None) + monkeypatch.setattr( + transform_cls, + "execute", + lambda self, config, cluster, expname, run_after=None: received.update(run_after=run_after), + ) + + runner = WorkflowRunner(str(GRPO_WORKFLOW)) + stages = ["validate_questions", "data_transformation"] + runner._run_stage("validate_questions", environment=environment, stages_to_run=stages) + runner._run_stage("data_transformation", environment=environment, stages_to_run=stages) + + submitted = {expname for expname, _ in submissions} + chained = {name for _, run_after in submissions for name in run_after} + assert sorted(received["run_after"]) == sorted(submitted - chained)