fix (cli): Correct the ordering of steps in multi-step experiments

mitosis/__init__.py

@@ -73,7 +73,7 @@ def variant_types():
 
 def load_trial_data(hexstr: str, *, trials_folder: Optional[Path | str] = None):
     trial = locate_trial_folder(hexstr, trials_folder=trials_folder)
-    all_files = glob("results*.dill", root_dir=trial)
+    all_files = sorted(glob("results*.dill", root_dir=trial))
     results = []
     for file in all_files:
         with open(trial / file, "rb") as fh:

mitosis/__main__.py

@@ -103,15 +103,24 @@ def _split_param_str(paramstr: str) -> CLIParam:
     return CLIParam(ex_step, track, arg_name, var_name)
 
 
+def _lookup_step_names(
+    experiment: list[str], config: str
+) -> dict[str, tuple[str, str]]:
+    proj_steps = _disk.load_mitosis_steps(config)
+    steps = {}
+    for step_name in experiment:
+        steps[step_name] = proj_steps[step_name]
+    return steps
+
+
 def _process_cl_args(args: argparse.Namespace) -> dict[str, Any]:
     ep_tups = [_split_param_str(epstr) for epstr in args.eval_param]
     lp_tups = [_split_param_str(lpstr) for lpstr in args.param]
 
     if args.experiment:
         if args.module:
             raise RuntimeError("Cannot use -m option if also passing experiment steps.")
-        proj_steps = _disk.load_mitosis_steps(args.config)
-        steps = dict(filter(lambda k: k[0] in args.experiment, proj_steps.items()))
+        steps = _lookup_step_names(args.experiment, args.config)
     elif args.module:
         mod: str = args.module
         steps = normalize_modinput(args.module)

mitosis/_disk.py

@@ -30,7 +30,7 @@ def load_mitosis_steps(
         )
     if any(not isinstance(vals, list) or len(vals) != 2 for vals in result.values()):
         raise RuntimeError("tool.mitosis.steps table is malformed")
-    return {k: tuple(v) for k, v in result.items()}
+    return {k: tuple(v)[:2] for k, v in result.items()}
 
 
 @lru_cache

mitosis/tests/test_all.py

@@ -1,3 +1,4 @@
+import pickle
 import sys
 from pathlib import Path
 from types import ModuleType
@@ -139,11 +140,25 @@ def test_mock_experiment(mock_steps, tmp_path):
         trials_folder=tmp_path,
     )
     data = mitosis.load_trial_data(exp_key, trials_folder=tmp_path)
-    assert len(data[1]["data"]) == 5
+    assert len(data[0]["data"]) == 5
     metadata = mitosis._disk.locate_trial_folder(exp_key, trials_folder=tmp_path)
     assert (metadata / "experiment").resolve().exists()
 
 
+def test_load_results_order(tmp_path):
+    exp_key = "test_results"
+    (tmp_path / exp_key).mkdir()
+    filename1 = "results1.dill"
+    filename2 = "results0.dill"
+    with open(tmp_path / exp_key / filename1, "wb") as f1:
+        pickle.dump("second", f1)
+    with open(tmp_path / exp_key / filename2, "wb") as f2:
+        pickle.dump("first", f2)
+    result = mitosis.load_trial_data(exp_key, trials_folder=tmp_path)
+    expected = ["first", "second"]
+    assert result == expected
+
+
 def test_mod_metadata_debug(mock_steps, tmp_path):
     hexstr = mitosis.run(
         mock_steps,

mitosis/tests/test_cli.py

@@ -8,6 +8,7 @@
 import pytest
 
 from mitosis.__main__ import _create_parser
+from mitosis.__main__ import _lookup_step_names
 from mitosis.__main__ import _process_cl_args
 from mitosis.__main__ import _split_param_str
 from mitosis.__main__ import main
@@ -18,6 +19,30 @@
 from mitosis.tests.mock_part1 import Klass
 
 
+def test_cli_step_ordering():
+    # GH 48
+    expected = ["fit_eval", "data"]
+    args = Namespace(
+        experiment=expected,
+        module=None,
+        debug=True,
+        config="mitosis/tests/test_pyproject.toml",
+        trials_folder=None,
+        eval_param=[],
+        param=[],
+    )
+    result = [step.name for step in _process_cl_args(args)["steps"]]
+    assert result == expected
+
+
+def test_lookup_step_names_ordering():
+    # GH 48
+    expected = ["fit_eval", "data"]
+    steps = _lookup_step_names(expected, "mitosis/tests/test_pyproject.toml")
+    result = list(steps.keys())
+    assert result == expected
+
+
 @pytest.mark.parametrize(
     ("params", "eval_params"),
     (([], []), (["foo=test"], []), ([], ["seed=1"])),