Merge pull request #2265 from transformerlab/add/task-queue-provider-flag

Add --provider flag to lab task queue

Author: deep1401

cli/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "transformerlab-cli"
-version = "0.0.64"
+version = "0.0.65"
 description = "Transformer Lab CLI"
 requires-python = ">=3.10"
 authors = [{ name = "Transformer Lab", email = "hello@transformerlab.ai" }]

cli/src/transformerlab_cli/commands/task.py

@@ -1045,6 +1045,20 @@ def _prompt_provider(providers: list[dict]) -> dict:
             console.print("[error]Please enter a valid number[/error]")
 
 
+def _resolve_provider(providers: list[dict], identifier: str) -> dict:
+    """Resolve a provider by name (case-insensitive) or exact id.
+
+    Raises typer.BadParameter listing available providers if no match.
+    """
+    match = next((p for p in providers if str(p.get("name", "")).lower() == identifier.lower()), None)
+    if match is None:
+        match = next((p for p in providers if str(p.get("id", "")) == identifier), None)
+    if match is None:
+        available = ", ".join(p.get("name", p.get("id", "?")) for p in providers) or "(none)"
+        raise typer.BadParameter(f"No provider matching {identifier!r}. Available providers: {available}")
+    return match
+
+
 def _prompt_parameters(parameters: dict) -> dict:
     """Prompt user for each parameter value, showing defaults."""
     if not parameters:
@@ -1112,6 +1126,7 @@ def queue_task(
     enable_profiling_torch: bool = False,
     use_spot: bool = False,
     image: str | None = None,
+    provider_name: str | None = None,
 ) -> None:
     """Queue a task on a compute provider."""
     with console.status("[bold success]Fetching task...[/bold success]", spinner="dots"):
@@ -1137,7 +1152,10 @@ def queue_task(
         console.print("[error]Error:[/error] No compute providers available. Add one in team settings first.")
         raise typer.Exit(1)
 
-    if interactive:
+    if provider_name:
+        provider = _resolve_provider(providers, provider_name)
+        console.print(f"[dim]Using provider: {provider.get('name')}[/dim]")
+    elif interactive:
         provider = _prompt_provider(providers)
     else:
         task_provider_id = task.get("provider_id")
@@ -1199,6 +1217,14 @@ def queue_task(
 def command_task_queue(
     task_id: str = typer.Argument(..., help="Task ID to queue"),
     no_interactive: bool = typer.Option(False, "--no-interactive", help="Skip interactive prompts, use defaults"),
+    provider: str | None = typer.Option(
+        None,
+        "--provider",
+        help=(
+            "Provider name (as shown in `lab provider list`) or id to run on. Overrides the task's "
+            "compute_provider and skips the interactive provider prompt. Errors if no provider matches."
+        ),
+    ),
     description: str | None = typer.Option(
         None,
         "--description",
@@ -1265,6 +1291,7 @@ def command_task_queue(
         enable_profiling_torch=enable_profiling_torch,
         use_spot=spot,
         image=image,
+        provider_name=provider,
     )
 
 

cli/tests/commands/test_task.py

@@ -231,6 +231,42 @@ def test_task_queue_no_interactive_picks_is_default_provider(_mock_exp, _mock_ge
     assert "p_default" in path, f"expected launch on default provider, got path: {path}"
 
 
+@patch("transformerlab_cli.commands.task.api.post_json", return_value=_mock_resp({"job_id": "j1"}))
+@patch("transformerlab_cli.commands.task.fetch_providers", return_value=SAMPLE_PROVIDERS_MULTI)
+@patch("transformerlab_cli.commands.task.api.get", return_value=_mock_resp(SAMPLE_TASK))
+@patch("transformerlab_cli.util.config.require_current_experiment", return_value="exp1")
+def test_task_queue_provider_flag_selects_by_name(_mock_exp, _mock_get, _mock_providers, mock_post):
+    """`--provider <name>` (case-insensitive) overrides the default and launches on that provider."""
+    result = runner.invoke(app, ["task", "queue", "t1", "--no-interactive", "--provider", "runpod", "-m", "x"])
+    assert result.exit_code == 0, result.output
+    path, body = mock_post.call_args.args
+    assert "p_other" in path, f"expected launch on named provider, got path: {path}"
+    assert body["provider_name"] == "Runpod"
+
+
+@patch("transformerlab_cli.commands.task.api.post_json", return_value=_mock_resp({"job_id": "j1"}))
+@patch("transformerlab_cli.commands.task.fetch_providers", return_value=SAMPLE_PROVIDERS_MULTI)
+@patch("transformerlab_cli.commands.task.api.get", return_value=_mock_resp(SAMPLE_TASK))
+@patch("transformerlab_cli.util.config.require_current_experiment", return_value="exp1")
+def test_task_queue_provider_flag_skips_prompt_in_interactive(_mock_exp, _mock_get, _mock_providers, mock_post):
+    """`--provider` resolves without prompting even when interactive (no --no-interactive)."""
+    result = runner.invoke(app, ["task", "queue", "t1", "--provider", "p_other", "-m", "x"], input="y\n")
+    assert result.exit_code == 0, result.output
+    path, _body = mock_post.call_args.args
+    assert "p_other" in path, f"expected launch on named provider, got path: {path}"
+
+
+@patch("transformerlab_cli.commands.task.api.post_json", return_value=_mock_resp({"job_id": "j1"}))
+@patch("transformerlab_cli.commands.task.fetch_providers", return_value=SAMPLE_PROVIDERS_MULTI)
+@patch("transformerlab_cli.commands.task.api.get", return_value=_mock_resp(SAMPLE_TASK))
+@patch("transformerlab_cli.util.config.require_current_experiment", return_value="exp1")
+def test_task_queue_provider_flag_unknown_errors(_mock_exp, _mock_get, _mock_providers, mock_post):
+    """`--provider <unknown>` exits non-zero and does not launch anything."""
+    result = runner.invoke(app, ["task", "queue", "t1", "--no-interactive", "--provider", "nope", "-m", "x"])
+    assert result.exit_code != 0
+    mock_post.assert_not_called()
+
+
 SAMPLE_TASK_WITH_PARAMS = {
     "id": "t1",
     "name": "finetune",