Merge branch 'main' into fix/new-playwright

Author: deep1401

cli/pyproject.toml

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

cli/src/transformerlab_cli/commands/task.py

@@ -825,6 +825,8 @@ def build_launch_payload(
     param_values: dict | None = None,
     resource_overrides: dict | None = None,
     description: str | None = None,
+    enable_profiling: bool = False,
+    enable_profiling_torch: bool = False,
 ) -> dict:
     """Build the payload for launching a task on a provider."""
     cfg = task.get("config") or {}
@@ -852,6 +854,8 @@ def pick(field: str):
         "accelerators": pick("accelerators"),
         "num_nodes": pick("num_nodes"),
         "minutes_requested": pick("minutes_requested"),
+        "enable_profiling": enable_profiling,
+        "enable_profiling_torch": enable_profiling_torch,
         "env_vars": task.get("env_vars", {}),
         "parameters": task.get("parameters", {}),
         "config": param_values if param_values else None,
@@ -1016,6 +1020,8 @@ def queue_task(
     interactive: bool = True,
     description: str | None = None,
     param_overrides: dict | None = None,
+    enable_profiling: bool = False,
+    enable_profiling_torch: bool = False,
 ) -> None:
     """Queue a task on a compute provider."""
     with console.status("[bold success]Fetching task...[/bold success]", spinner="dots"):
@@ -1067,8 +1073,17 @@ def queue_task(
         param_values = {k: (v.get("default", "") if isinstance(v, dict) else v) for k, v in parameters.items()}
     param_values.update(overrides)
 
+    if enable_profiling_torch and not enable_profiling:
+        raise typer.BadParameter("--enable-profiling-torch requires --enable-profiling.")
+
     payload = build_launch_payload(
-        task, provider.get("name"), param_values, resource_overrides, description=description
+        task,
+        provider.get("name"),
+        param_values,
+        resource_overrides,
+        description=description,
+        enable_profiling=enable_profiling,
+        enable_profiling_torch=enable_profiling_torch,
     )
     provider_id = provider.get("id")
 
@@ -1105,6 +1120,16 @@ def command_task_queue(
             "scalar (e.g. score=0.42 -> float, enabled=true -> bool). Unknown keys fail."
         ),
     ),
+    enable_profiling: bool = typer.Option(
+        False,
+        "--enable-profiling",
+        help="Enable system profiling for this run (CPU/GPU/memory sampling).",
+    ),
+    enable_profiling_torch: bool = typer.Option(
+        False,
+        "--enable-profiling-torch",
+        help="Enable torch profiler trace export for this run (requires --enable-profiling).",
+    ),
 ):
     """Queue a task on a compute provider."""
     current_experiment = require_current_experiment()
@@ -1119,6 +1144,8 @@ def command_task_queue(
         interactive=not no_interactive,
         description=description,
         param_overrides=param_overrides,
+        enable_profiling=enable_profiling,
+        enable_profiling_torch=enable_profiling_torch,
     )
 
 

cli/tests/commands/test_task.py

@@ -95,6 +95,14 @@ def test_build_launch_payload_omits_description_by_default():
     assert payload["description"] is None
 
 
+def test_build_launch_payload_includes_profiling_flags():
+    """build_launch_payload forwards profiling flags to the launch API body."""
+    task = {"id": "t1", "name": "finetune", "experiment_id": "exp1", "run": "python main.py"}
+    payload = build_launch_payload(task, "Local", enable_profiling=True, enable_profiling_torch=True)
+    assert payload["enable_profiling"] is True
+    assert payload["enable_profiling_torch"] is True
+
+
 SAMPLE_TASK = {
     "id": "t1",
     "name": "finetune",
@@ -264,6 +272,44 @@ def test_task_queue_param_when_task_has_no_parameters_errors(_mock_exp, _mock_ge
     mock_post.assert_not_called()
 
 
+@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)
+@patch("transformerlab_cli.commands.task.api.get", return_value=_mock_resp(SAMPLE_TASK))
+@patch("transformerlab_cli.commands.task.require_current_experiment", return_value="exp1")
+def test_task_queue_enable_profiling_flags(_mock_exp, _mock_get, _mock_providers, mock_post):
+    """Queue command forwards profiling flags into launch payload."""
+    result = runner.invoke(
+        app,
+        [
+            "task",
+            "queue",
+            "t1",
+            "--no-interactive",
+            "--enable-profiling",
+            "--enable-profiling-torch",
+        ],
+    )
+    assert result.exit_code == 0, result.output
+    _path, body = mock_post.call_args.args
+    assert body["enable_profiling"] is True
+    assert body["enable_profiling_torch"] is True
+
+
+@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)
+@patch("transformerlab_cli.commands.task.api.get", return_value=_mock_resp(SAMPLE_TASK))
+@patch("transformerlab_cli.commands.task.require_current_experiment", return_value="exp1")
+def test_task_queue_enable_torch_profiling_requires_profiling(_mock_exp, _mock_get, _mock_providers, mock_post):
+    """Torch profiling flag without base profiling fails with a clear error."""
+    result = runner.invoke(
+        app,
+        ["task", "queue", "t1", "--no-interactive", "--enable-profiling-torch"],
+    )
+    assert result.exit_code != 0
+    assert "requires --enable-profiling" in strip_ansi(result.output).lower()
+    mock_post.assert_not_called()
+
+
 @patch(
     "transformerlab_cli.commands.task.api.post_json",
     side_effect=[