Repository: mmcdermott/MEDS_transforms

Issue:
[Proposal] Consider having main entry run point explicitly take both the pipeline config yaml (or package) and the stage name, and register stages without a pipeline.
> So, one additional challenge with this is that the intention for the generic stages in the MEDS-Transforms repo (or even in derived packages) is that they could be used in further packages on their own. So, registering a stage in a pipeline specific manner is therefore an issue. This can be solved by moving the final hydra main construction to the actual main entry point, rather than at registration time, so it can be done in a pipeline specific manner, rather than in a stage specific manner. This is likely a general improvement in design, though could be done independently of this bigger change. 

 _Originally posted by @mmcdermott in [#283](https://github.com/mmcdermott/MEDS_transforms/issues/283#issuecomment-2761672143)_

Discussion:

Patch:
diff --git a/src/MEDS_transforms/__main__.py b/src/MEDS_transforms/__main__.py
index f01a619..6c60601 100644
--- a/src/MEDS_transforms/__main__.py
+++ b/src/MEDS_transforms/__main__.py
@@ -1,8 +1,18 @@
+import os
 import sys
 from importlib.metadata import entry_points
+from importlib.resources import files
+from pathlib import Path
+
+import hydra
+from omegaconf import OmegaConf
 
 from . import __package_name__
 
+HELP_STRS = {"--help", "-h", "help", "h"}
+PKG_PFX = "pkg://"
+YAML_EXTENSIONS = {"yaml", "yml"}
+
 
 def get_all_stages():
     """Get all available stages."""
@@ -10,26 +20,81 @@ def get_all_stages():
     return {name: eps[name] for name in eps.names}
 
 
+def print_help_stage():
+    """Print help for all stages."""
+    print(f"Usage: {sys.argv[0]} <pipeline_yaml> <stage_name> [args]")
+    print(
+        "  * pipeline_yaml: Path to the pipeline YAML file on disk or in the "
+        f"'{PKG_PFX}<pkg_name>.<relative_path>' format."
+    )
+    print("  * stage_name: Name of the stage to run.")
+    print()
+    print("Available stages:")
+    all_stages = get_all_stages()
+    for name in sorted(all_stages):
+        print(f"  - {name}")
+
+
+def resolve_pipeline_yaml(pipeline_yaml: str):
+    """Resolve the pipeline YAML file path."""
+    if pipeline_yaml.startswith(PKG_PFX):
+        pipeline_yaml = pipeline_yaml[len(PKG_PFX) :]
+        pipeline_parts = pipeline_yaml.split(".")
+
+        if pipeline_parts[-1] not in YAML_EXTENSIONS:
+            raise ValueError(
+                f"Invalid pipeline YAML path '{pipeline_yaml}'. "
+                f"Expected a file with one of the following extensions: {YAML_EXTENSIONS}"
+            )
+
+        pkg_name = pipeline_parts[0]
+        suffix = pipeline_parts[-1]
+        relative_path = Path(os.path.join(*pipeline_parts[1:-1])).with_suffix(f".{suffix}")
+
+        try:
+            pipeline_yaml = files(pkg_name) / relative_path
+        except ImportError:
+            raise ValueError(f"Package '{pkg_name}' not found. Please check the package name.")
+    else:
+        pipeline_yaml = Path(pipeline_yaml)
+
+    if not pipeline_yaml.is_file():
+        raise FileNotFoundError(f"Pipeline YAML file '{pipeline_yaml}' does not exist.")
+
+    return pipeline_yaml
+
+
 def run_stage():
     """Run a stage based on command line arguments."""
 
     all_stages = get_all_stages()
 
-    if len(sys.argv) < 2 or sys.argv[1] in ("--help", "-h"):
-        print(f"Usage: {sys.argv[0]} <stage_name> [args]")
-        print("Available stages:")
-        for name in sorted(all_stages):
-            print(f"  - {name}")
-        if len(sys.argv) < 2:
-            sys.exit(1)
-        else:
-            sys.exit(0)
+    if len(sys.argv) < 2:
+        print_help_stage()
+        sys.exit(1)
+    elif sys.argv[1] in HELP_STRS:
+        print_help_stage()
+        sys.exit(0)
+    elif len(sys.argv) < 3:
+        print_help_stage()
+        sys.exit(1)
 
-    stage_name = sys.argv[1]
-    sys.argv = sys.argv[1:]  # remove dispatcher argument
+    pipeline_yaml = resolve_pipeline_yaml(sys.argv[1])
+    stage_name = sys.argv[2]
+    sys.argv = sys.argv[2:]  # remove dispatcher arguments
 
     if stage_name not in all_stages:
         raise ValueError(f"Stage '{stage_name}' not found.")
 
     main_fn = all_stages[stage_name].load()
-    main_fn()
+
+    hydra_wrapper = hydra.main(
+        version_base=None,
+        config_path=str(pipeline_yaml.parent),
+        config_name=pipeline_yaml.stem,
+    )
+
+    OmegaConf.register_new_resolver("stage_name", lambda: stage_name)
+    OmegaConf.register_new_resolver("stage_docstring", lambda: main_fn.__doc__.replace("$", "$$"))
+
+    hydra_wrapper(main_fn)()
diff --git a/src/MEDS_transforms/runner.py b/src/MEDS_transforms/runner.py
index 00aadc8..7d3c7f3 100644
--- a/src/MEDS_transforms/runner.py
+++ b/src/MEDS_transforms/runner.py
@@ -155,8 +155,11 @@ def run_stage(
         ...     },
         ... })
         >>> run_stage(cfg, "reshard_to_split", runner_fn=fake_shell_succeed)
-        MEDS_transform-stage reshard_to_split --config-dir=... --config-name=pipeline_config
-            'hydra.searchpath=[pkg://MEDS_transforms.configs]' stage=reshard_to_split
+        MEDS_transform-stage pipeline_config.yaml reshard_to_split
+            --config-dir=...
+            --config-name=pipeline_config
+            'hydra.searchpath=[pkg://MEDS_transforms.configs]'
+            stage=reshard_to_split
         >>> run_stage(
         ...     cfg, "fit_vocabulary_indices", runner_fn=fake_shell_succeed
         ... )
@@ -195,7 +198,7 @@ def run_stage(
     elif "_script" in stage_config:
         script = stage_config._script
     else:
-        script = f"MEDS_transform-stage {stage_name}"
+        script = f"MEDS_transform-stage {str(pipeline_config_fp)} {stage_name}"
 
     command_parts = [
         script,
diff --git a/src/MEDS_transforms/stage.py b/src/MEDS_transforms/stage.py
index d5f2d1f..bf538fd 100644
--- a/src/MEDS_transforms/stage.py
+++ b/src/MEDS_transforms/stage.py
@@ -5,12 +5,9 @@ import inspect
 import logging
 from collections.abc import Callable
 from enum import StrEnum
-from pathlib import Path
 
-import hydra
-from omegaconf import DictConfig, OmegaConf
+from omegaconf import DictConfig
 
-from .configs import PREPROCESS_CONFIG_YAML
 from .mapreduce import ANY_COMPUTE_FN_T, map_stage, mapreduce_stage
 
 logger = logging.getLogger(__name__)
@@ -92,7 +89,6 @@ def get_stage_main(
     main_fn: MAIN_FN_T | None = None,
     map_fn: ANY_COMPUTE_FN_T | None = None,
     reduce_fn: ANY_COMPUTE_FN_T | None = None,
-    config_path: Path | None = None,
     stage_name: str | None = None,
     stage_docstring: str | None = None,
 ) -> MAIN_FN_T:
@@ -100,15 +96,6 @@ def get_stage_main(
 
     mode = StageType.from_fns(main_fn, map_fn, reduce_fn)
 
-    if config_path is None:
-        config_path = PREPROCESS_CONFIG_YAML
-
-    hydra_wrapper = hydra.main(
-        version_base=None,
-        config_path=str(config_path.parent),
-        config_name=config_path.stem,
-    )
-
     if stage_name is None:
         stage_name = (main_fn or map_fn).__module__.split(".")[-1]
 
@@ -123,7 +110,6 @@ def get_stage_main(
                 return map_stage(cfg, map_fn)
 
             main_fn.__name__ = stage_name
-            main_fn.__doc__ = stage_docstring
 
         case StageType.MAPREDUCE:
             stage_docstring = stage_docstring or (
@@ -135,17 +121,10 @@ def get_stage_main(
                 return mapreduce_stage(cfg, map_fn, reduce_fn)
 
             main_fn.__name__ = stage_name
-            main_fn.__doc__ = stage_docstring
-
-    hydra_wraped_main = hydra_wrapper(main_fn)
 
-    @functools.wraps(hydra_wraped_main)
-    def wrapped_main(*args, **kwargs):
-        OmegaConf.register_new_resolver("stage_name", lambda: stage_name)
-        OmegaConf.register_new_resolver("stage_docstring", lambda: stage_docstring.replace("$", "$$"))
-        return hydra_wraped_main(*args, **kwargs)
+    main_fn.__doc__ = stage_docstring
 
-    return wrapped_main
+    return main_fn
 
 
 def MEDS_transforms_stage(*args, **kwargs):
@@ -186,7 +165,7 @@ def MEDS_transforms_stage(*args, **kwargs):
 
     Args:
         *args: Positional arguments. If used as a decorator, this should be a single function.
-        **kwargs: Keyword arguments. These can include `main_fn`, `map_fn`, `reduce_fn`, `config_path`,
+        **kwargs: Keyword arguments. These can include `main_fn`, `map_fn`, `reduce_fn`,
             `stage_name`, and `stage_docstring`. Other keyword arguments will cause an error. Not all keyword
             arguments are required for all usages of the decorator.
 
@@ -300,7 +279,7 @@ def MEDS_transforms_stage(*args, **kwargs):
         Keyword arguments can be used to specify the config path, stage name, and docstring. These are used to
         modify the returned stage name and docstring, and the config path is used to set up the hydra wrapper:
 
-        >>> @MEDS_transforms_stage(config_path=Path("f.yaml"), stage_name="bar", stage_docstring="baz")
+        >>> @MEDS_transforms_stage(stage_name="bar", stage_docstring="baz")
         ... def special_map(cfg: DictConfig):
         ...     '''Not baz'''
         ...     return "map"
@@ -311,7 +290,7 @@ def MEDS_transforms_stage(*args, **kwargs):
 
         If specified with no positional arguments and only non-determining keyword arguments (meaning no main
         or map function), then a decorator is returned that itself takes a function:
-        >>> fntr = MEDS_transforms_stage(config_path=Path("f.yaml"), stage_name="bar", stage_docstring="baz")
+        >>> fntr = MEDS_transforms_stage(stage_name="bar", stage_docstring="baz")
         >>> fn = fntr(main)
         >>> hasattr(fn, "main")
         False


Test patch:
diff --git a/tests/__init__.py b/tests/__init__.py
index dca4b15..f65f5e4 100644
--- a/tests/__init__.py
+++ b/tests/__init__.py
@@ -2,7 +2,7 @@
 RUNNER_SCRIPT = "MEDS_transform-pipeline"
 
 # Stages
-__stage_pattern = "MEDS_transform-stage {stage_name}"
+__stage_pattern = "MEDS_transform-stage pkg://MEDS_transforms.configs._preprocess.yaml {stage_name}"
 
 AGGREGATE_CODE_METADATA_SCRIPT = __stage_pattern.format(stage_name="aggregate_code_metadata")
 FIT_VOCABULARY_INDICES_SCRIPT = __stage_pattern.format(stage_name="fit_vocabulary_indices")
diff --git a/tests/test_main.py b/tests/test_main.py
index 0508686..f7e7050 100644
--- a/tests/test_main.py
+++ b/tests/test_main.py
@@ -1,5 +1,9 @@
+import re
 import subprocess
 
+SCRIPT_TEMPLATE = "MEDS_transform-stage {pipeline} {stage_name}"
+PKG_BASE = "pkg://MEDS_transforms.configs.{config}"
+
 
 def test_stage_entry_point_help():
     result = subprocess.run("MEDS_transform-stage", check=False, shell=True, capture_output=True)
@@ -12,10 +16,30 @@ def test_stage_entry_point_help():
     assert result.returncode == 0
     assert result.stdout.decode() == help_str
 
-
-def test_stage_entry_point_invalid():
-    result = subprocess.run(
-        "MEDS_transform-stage invalid_stage", check=False, shell=True, capture_output=True
-    )
+    result = subprocess.run("MEDS_transform-stage foo", check=False, shell=True, capture_output=True)
     assert result.returncode != 0
-    assert "Stage 'invalid_stage' not found." in result.stderr.decode()
+    assert result.stdout.decode() == help_str
+
+
+def test_stage_entry_point_errors():
+    for pipeline, stage, want_err in [
+        ("not_real.yaml", "occlude_outliers", "Pipeline YAML file 'not_real.yaml' does not exist."),
+        (PKG_BASE.format(config="_preprocess.yaml"), "not_real_stage", "Stage 'not_real_stage' not found."),
+        (
+            "pkg://pkg.bad_suffix.json",
+            "occlude_outliers",
+            re.compile("Invalid pipeline YAML path .* Expected a file with one of the following extensions"),
+        ),
+        (
+            "pkg://non_existent_pkg.file.yaml",
+            "occlude_outliers",
+            re.compile("Package 'non_existent_pkg' not found"),
+        ),
+    ]:
+        script = SCRIPT_TEMPLATE.format(pipeline=pipeline, stage_name=stage)
+        result = subprocess.run(script, check=False, shell=True, capture_output=True)
+        assert result.returncode != 0
+        if isinstance(want_err, str):
+            assert want_err in result.stderr.decode()
+        else:
+            assert re.search(want_err, result.stderr.decode()) is not None