Fix run_after names for stages with custom experiment naming

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`

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
 

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]

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.
 

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"]

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"):