Add alias functionality for train, launch, evaluate, and infer in the oumi CLI (#1618)

Author: taenin

src/oumi/cli/alias.py

@@ -0,0 +1,72 @@
+# Copyright 2025 - Oumi
+#
+# 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.
+
+from enum import Enum
+
+from oumi.utils.logging import logger
+
+
+class AliasType(str, Enum):
+    """The type of configs we support with aliases."""
+
+    TRAIN = "train"
+    EVAL = "eval"
+    INFER = "infer"
+    JOB = "job"
+
+
+_ALIASES: dict[str, dict[AliasType, str]] = {
+    "llama4-scout": {
+        AliasType.TRAIN: "oumi://configs/recipes/llama4/sft/scout_base_full/train.yaml",
+        AliasType.JOB: "oumi://configs/recipes/llama4/sft/scout_base_full/train.yaml",
+    },
+    "llama4-scout-instruct-lora": {
+        AliasType.TRAIN: "oumi://configs/recipes/llama4/sft/scout_instruct_lora/train.yaml",
+    },
+    "llama4-scout-instruct-qlora": {
+        AliasType.TRAIN: "oumi://configs/recipes/llama4/sft/scout_instruct_qlora/train.yaml",
+    },
+    "llama4-scout-instruct": {
+        AliasType.TRAIN: "oumi://configs/recipes/llama4/sft/scout_instruct_full/train.yaml",
+        AliasType.INFER: "oumi://configs/recipes/llama4/inference/scout_instruct_infer.yaml",
+        AliasType.JOB: "oumi://configs/recipes/llama4/sft/scout_instruct_full/gcp_job.yaml",
+        AliasType.EVAL: "oumi://configs/recipes/llama4/evaluation/scout_instruct_eval.yaml",
+    },
+    "llama4-maverick": {
+        AliasType.INFER: "oumi://configs/recipes/llama4/inference/maverick_instruct_together_infer.yaml",
+    },
+}
+
+
+def try_get_config_name_for_alias(
+    alias: str,
+    alias_type: AliasType,
+) -> str:
+    """Gets the config path for a given alias.
+
+    This function resolves the config path for a given alias and alias type.
+    If the alias is not found, the original alias is returned.
+
+    Args:
+        alias (str): The alias to resolve.
+        alias_type (AliasType): The type of config to resolve.
+
+    Returns:
+        str: The resolved config path (or the original alias if not found).
+    """
+    if alias in _ALIASES and alias_type in _ALIASES[alias]:
+        config_path = _ALIASES[alias][alias_type]
+        logger.info(f"Resolved alias '{alias}' to '{config_path}'")
+        return config_path
+    return alias

src/oumi/cli/evaluate.py

@@ -18,6 +18,7 @@
 from rich.table import Table
 
 import oumi.cli.cli_utils as cli_utils
+from oumi.cli.alias import AliasType, try_get_config_name_for_alias
 from oumi.utils.logging import logger
 
 
@@ -42,7 +43,7 @@ def evaluate(
 
     config = str(
         cli_utils.resolve_and_fetch_config(
-            config,
+            try_get_config_name_for_alias(config, AliasType.EVAL),
         )
     )
     with cli_utils.CONSOLE.status(

src/oumi/cli/infer.py

@@ -19,6 +19,7 @@
 from rich.table import Table
 
 import oumi.cli.cli_utils as cli_utils
+from oumi.cli.alias import AliasType, try_get_config_name_for_alias
 from oumi.utils.logging import logger
 
 _DEFAULT_CLI_PDF_DPI: Final[int] = 200
@@ -79,7 +80,7 @@ def infer(
 
     config = str(
         cli_utils.resolve_and_fetch_config(
-            config,
+            try_get_config_name_for_alias(config, AliasType.INFER),
         )
     )
 

src/oumi/cli/launch.py

@@ -25,6 +25,7 @@
 from rich.text import Text
 
 import oumi.cli.cli_utils as cli_utils
+from oumi.cli.alias import AliasType, try_get_config_name_for_alias
 from oumi.utils.git_utils import get_git_root_dir
 from oumi.utils.logging import logger
 from oumi.utils.version_utils import is_dev_build
@@ -324,7 +325,7 @@ def run(
 
     config = str(
         cli_utils.resolve_and_fetch_config(
-            config,
+            try_get_config_name_for_alias(config, AliasType.JOB),
         )
     )
 
@@ -490,7 +491,7 @@ def up(
 
     config = str(
         cli_utils.resolve_and_fetch_config(
-            config,
+            try_get_config_name_for_alias(config, AliasType.JOB),
         )
     )
 

src/oumi/cli/train.py

@@ -17,6 +17,7 @@
 import typer
 
 import oumi.cli.cli_utils as cli_utils
+from oumi.cli.alias import AliasType, try_get_config_name_for_alias
 from oumi.utils.logging import logger
 
 
@@ -41,7 +42,7 @@ def train(
 
     config = str(
         cli_utils.resolve_and_fetch_config(
-            config,
+            try_get_config_name_for_alias(config, AliasType.TRAIN),
         )
     )
     with cli_utils.CONSOLE.status(

tests/unit/cli/test_cli_alias.py

@@ -0,0 +1,28 @@
+from oumi.cli.alias import _ALIASES, AliasType, try_get_config_name_for_alias
+
+
+def test_alias_all_entries():
+    for alias in _ALIASES:
+        for alias_type in _ALIASES[alias]:
+            config_path = try_get_config_name_for_alias(alias, alias_type)
+            assert config_path == _ALIASES[alias][alias_type], (
+                f"Alias '{alias}' with type '{alias_type}' did not resolve correctly."
+                f" Expected: {config_path}, Actual: {_ALIASES[alias][alias_type]}"
+            )
+
+
+def test_alias_not_found():
+    alias = "non_existent_alias"
+    alias_type = AliasType.TRAIN
+    config_path = try_get_config_name_for_alias(alias, alias_type)
+    assert (
+        config_path == alias
+    ), f"Expected the original alias '{alias}' to be returned."
+
+
+def test_alias_type_not_found():
+    alias = "llama4-scout"
+    config_path = try_get_config_name_for_alias(alias, AliasType.EVAL)
+    assert (
+        config_path == alias
+    ), f"Expected the original alias '{alias}' to be returned."

tests/unit/cli/test_cli_evaluate.py

@@ -8,6 +8,7 @@
 from typer.testing import CliRunner
 
 import oumi
+import oumi.cli.alias
 from oumi.cli.cli_utils import CONTEXT_ALLOW_EXTRA_ARGS
 from oumi.cli.evaluate import evaluate
 from oumi.core.configs import (
@@ -26,6 +27,12 @@ def mock_fetch():
         yield m_fetch
 
 
+@pytest.fixture
+def mock_alias():
+    with patch("oumi.cli.evaluate.try_get_config_name_for_alias") as try_alias:
+        yield try_alias
+
+
 def _create_eval_config() -> EvaluationConfig:
     return EvaluationConfig(
         output_dir="output/dir",
@@ -69,6 +76,21 @@ def test_evaluate_runs(app, mock_evaluate):
         mock_evaluate.assert_has_calls([call(config)])
 
 
+def test_evaluate_calls_alias(app, mock_evaluate, mock_alias):
+    with tempfile.TemporaryDirectory() as output_temp_dir:
+        yaml_path = str(Path(output_temp_dir) / "eval.yaml")
+        mock_alias.return_value = yaml_path
+        config: EvaluationConfig = _create_eval_config()
+        config.to_yaml(yaml_path)
+        _ = runner.invoke(app, ["--config", "an_alias"])
+        mock_alias.assert_has_calls(
+            [
+                call("an_alias", oumi.cli.alias.AliasType.EVAL),
+            ]
+        )
+        mock_evaluate.assert_has_calls([call(config)])
+
+
 def test_evaluate_unparsable_metrics(app, mock_evaluate):
     with tempfile.TemporaryDirectory() as output_temp_dir:
         mock_evaluate.return_value = [

tests/unit/cli/test_cli_infer.py

@@ -10,6 +10,7 @@
 from typer.testing import CliRunner
 
 import oumi
+from oumi.cli.alias import AliasType
 from oumi.cli.cli_utils import CONTEXT_ALLOW_EXTRA_ARGS
 from oumi.cli.infer import infer
 from oumi.core.configs import (
@@ -29,6 +30,12 @@ def mock_fetch():
         yield m_fetch
 
 
+@pytest.fixture
+def mock_alias():
+    with patch("oumi.cli.infer.try_get_config_name_for_alias") as try_alias:
+        yield try_alias
+
+
 def _create_inference_config() -> InferenceConfig:
     return InferenceConfig(
         model=ModelParams(
@@ -75,6 +82,19 @@ def test_infer_runs(app, mock_infer, mock_infer_interactive):
         )
 
 
+def test_infer_with_alias_runs(app, mock_infer, mock_infer_interactive, mock_alias):
+    with tempfile.TemporaryDirectory() as output_temp_dir:
+        yaml_path = str(Path(output_temp_dir) / "infer.yaml")
+        mock_alias.return_value = yaml_path
+        config: InferenceConfig = _create_inference_config()
+        config.to_yaml(yaml_path)
+        _ = runner.invoke(app, ["-i", "--config", "random_alias"])
+        mock_alias.assert_called_once_with("random_alias", AliasType.INFER)
+        mock_infer_interactive.assert_has_calls(
+            [call(config, input_image_bytes=None, system_prompt=None)]
+        )
+
+
 def test_infer_runs_interactive_by_default(app, mock_infer, mock_infer_interactive):
     with tempfile.TemporaryDirectory() as output_temp_dir:
         yaml_path = str(Path(output_temp_dir) / "infer.yaml")

tests/unit/cli/test_cli_launch.py

@@ -8,6 +8,7 @@
 from typer.testing import CliRunner
 
 import oumi
+from oumi.cli.alias import AliasType
 from oumi.cli.cli_utils import CONTEXT_ALLOW_EXTRA_ARGS
 from oumi.cli.launch import cancel, down, status, stop, up, which
 from oumi.cli.launch import run as launcher_run
@@ -32,6 +33,12 @@ def mock_fetch():
         yield m_fetch
 
 
+@pytest.fixture
+def mock_alias():
+    with patch("oumi.cli.launch.try_get_config_name_for_alias") as try_alias:
+        yield try_alias
+
+
 runner = CliRunner()
 
 
@@ -202,6 +209,52 @@ def test_launch_up_job(
         assert logger.level == logging.DEBUG
 
 
+def test_launch_up_job_with_alias(
+    app, mock_launcher, mock_pool, mock_version, mock_confirm, mock_fetch, mock_alias
+):
+    with tempfile.TemporaryDirectory() as output_temp_dir:
+        train_yaml_path = str(pathlib.Path(output_temp_dir) / "train.yaml")
+        config: TrainingConfig = _create_training_config()
+        config.to_yaml(train_yaml_path)
+        job_yaml_path = str(pathlib.Path(output_temp_dir) / "job.yaml")
+        mock_alias.return_value = job_yaml_path
+        job_config = _create_job_config(train_yaml_path)
+        job_config.to_yaml(job_yaml_path)
+        mock_launcher.JobConfig = JobConfig
+        mock_cluster = Mock()
+        job_status = JobStatus(
+            id="job_id",
+            cluster="cluster_id",
+            name="job_name",
+            status="running",
+            metadata="",
+            done=False,
+        )
+        mock_launcher.up.return_value = (mock_cluster, job_status)
+        mock_cluster.get_job.return_value = job_status = JobStatus(
+            id="job_id",
+            cluster="cluster_id",
+            name="job_name",
+            status="done",
+            metadata="",
+            done=True,
+        )
+        _ = runner.invoke(
+            app,
+            [
+                "up",
+                "--config",
+                "some_alias",
+                "--log-level",
+                "DEBUG",
+            ],
+        )
+        mock_fetch.assert_called_once_with(job_yaml_path)
+        mock_cluster.get_job.assert_has_calls([call("job_id")])
+        mock_alias.assert_called_once_with("some_alias", AliasType.JOB)
+        assert logger.level == logging.DEBUG
+
+
 def test_launch_up_job_dev_confirm(
     app, mock_launcher, mock_pool, mock_version, mock_confirm, mock_git_root, mock_fetch
 ):
@@ -560,6 +613,62 @@ def test_launch_run_job(
         assert logger.level == logging.CRITICAL
 
 
+def test_launch_run_job_with_alias(
+    app, mock_launcher, mock_pool, mock_version, mock_confirm, mock_fetch, mock_alias
+):
+    with tempfile.TemporaryDirectory() as output_temp_dir:
+        train_yaml_path = str(pathlib.Path(output_temp_dir) / "train.yaml")
+        config: TrainingConfig = _create_training_config()
+        config.to_yaml(train_yaml_path)
+        job_yaml_path = str(pathlib.Path(output_temp_dir) / "job.yaml")
+        mock_alias.return_value = job_yaml_path
+        job_config = _create_job_config(train_yaml_path)
+        job_config.to_yaml(job_yaml_path)
+        mock_launcher.JobConfig = JobConfig
+        mock_cluster = Mock()
+        job_status = JobStatus(
+            id="job_id",
+            cluster="cluster_id",
+            name="job_name",
+            status="running",
+            metadata="",
+            done=False,
+        )
+        mock_cloud = Mock()
+        mock_launcher.run.return_value = job_status
+        mock_launcher.get_cloud.side_effect = [mock_cloud, mock_cloud]
+        mock_cloud.get_cluster.side_effect = [mock_cluster, mock_cluster]
+        mock_cluster.get_job.return_value = job_status = JobStatus(
+            id="job_id",
+            cluster="cluster_id",
+            name="job_name",
+            status="done",
+            metadata="",
+            done=True,
+        )
+        _ = runner.invoke(
+            app,
+            [
+                "run",
+                "--config",
+                "some_alias",
+                "--cluster",
+                "cluster_id",
+                "-log",
+                "CRITICAL",
+            ],
+        )
+        mock_cluster.get_job.assert_has_calls([call("job_id"), call("job_id")])
+        mock_launcher.run.assert_called_once_with(job_config, "cluster_id")
+        mock_launcher.get_cloud.assert_has_calls([call("aws"), call("aws")])
+        mock_cloud.get_cluster.assert_has_calls(
+            [call("cluster_id"), call("cluster_id")]
+        )
+        mock_fetch.assert_called_once_with(job_yaml_path)
+        mock_alias.assert_called_once_with("some_alias", AliasType.JOB)
+        assert logger.level == logging.CRITICAL
+
+
 def test_launch_run_job_dev_confirm(
     app, mock_launcher, mock_pool, mock_version, mock_confirm, mock_git_root, mock_fetch
 ):