From 646a72bbf0776b5414088020579daebd7895dd21 Mon Sep 17 00:00:00 2001
From: Marc Romeyn <mromeijn@nvidia.com>
Date: Mon, 19 May 2025 16:10:53 +0200
Subject: [PATCH] Fix bug with a CLI overwrite

Signed-off-by: Marc Romeyn <mromeijn@nvidia.com>
---
 nemo_run/cli/lazy.py |  5 ++++-
 test/cli/test_api.py | 27 +++++++++++++++++++++++++++
 2 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/nemo_run/cli/lazy.py b/nemo_run/cli/lazy.py
index 882b018e..efee05b6 100644
--- a/nemo_run/cli/lazy.py
+++ b/nemo_run/cli/lazy.py
@@ -721,7 +721,10 @@ def _args_to_dictconfig(args: list[tuple[str, str, Any]]) -> DictConfig:
             if part not in current:
                 current[part] = {}
             elif not isinstance(current[part], dict):
-                current[part] = {}  # Convert to dict if it wasn't already
+                if isinstance(current[part], str):
+                    current[part] = {"_factory_": current[part]}
+                else:
+                    current[part] = {}  # Convert to dict if it wasn't already
             current = current[part]
 
         # Add the operator suffix if it's not a simple assignment
diff --git a/test/cli/test_api.py b/test/cli/test_api.py
index f0f32104..da359c12 100644
--- a/test/cli/test_api.py
+++ b/test/cli/test_api.py
@@ -631,6 +631,13 @@ def my_model(hidden_size: int = 256, num_layers: int = 3, activation: str = "rel
     return Model(hidden_size=hidden_size, num_layers=num_layers, activation=activation)
 
 
+@run.cli.factory
+@run.autoconvert
+def my_other_model(hidden_size: int = 512, num_layers: int = 3, activation: str = "relu") -> Model:
+    """Create a model configuration."""
+    return Model(hidden_size=hidden_size, num_layers=num_layers, activation=activation)
+
+
 @run.cli.factory
 def my_optimizer(
     learning_rate: float = 0.001, weight_decay: float = 1e-5, betas: List[float] = [0.9, 0.999]
@@ -880,6 +887,26 @@ class SomeObject:
         value_1: int
         value_2: int
 
+    def test_with_factory_and_overwrite(self, runner, app):
+        # Test CLI execution with factory and parameter overwrite
+        result = runner.invoke(
+            app,
+            [
+                "my_llm",
+                "train_model",
+                "model=my_other_model",
+                "model.num_layers=10",
+                "--yes",
+            ],
+            env={"INCLUDE_WORKSPACE_FILE": "false"},
+        )
+        assert result.exit_code == 0
+
+        output = result.stdout
+        assert "Training model with the following configuration:" in output
+        # Check that my_model_2's default hidden_size (512) is used
+        assert "Model: Model(hidden_size=512, num_layers=10, activation='relu')" in output
+
 
 class TestDefaultFactory:
     def test_default_factory(self):