Merge pull request #1927 from transformerlab/fix/stop-cluster-cli

Fix CLI to stop the cluster as well when a job stops

Author: deep1401

cli/src/transformerlab_cli/commands/job.py

@@ -31,6 +31,39 @@
 ACTIVE_JOB_STATUSES = {"RUNNING", "LAUNCHING", "INTERACTIVE", "WAITING"}
 
 
+def _extract_provider_cluster_from_job(job: dict) -> tuple[str | None, str | None]:
+    """Extract provider_id and cluster_name from a job payload."""
+    if not isinstance(job, dict):
+        return None, None
+    job_data = job.get("job_data", {})
+    if not isinstance(job_data, dict):
+        return None, None
+    provider_id = job_data.get("provider_id")
+    cluster_name = job_data.get("cluster_name")
+    if not provider_id or not cluster_name:
+        return None, None
+    return str(provider_id), str(cluster_name)
+
+
+def _stop_provider_cluster(experiment_id: str, job_id: str) -> None:
+    """Best-effort provider cluster stop to mirror GUI behavior."""
+    job_response = api.get(f"/experiment/{experiment_id}/jobs/{job_id}")
+    if job_response.status_code != 200:
+        return
+
+    provider_id, cluster_name = _extract_provider_cluster_from_job(job_response.json())
+    if not provider_id or not cluster_name:
+        return
+
+    cluster_stop_url = f"/compute_provider/providers/{provider_id}/clusters/{cluster_name}/stop?job_id={job_id}"
+    cluster_response = api.post(cluster_stop_url)
+    if cluster_response.status_code >= 400:
+        console.print(
+            f"[warning]Warning:[/warning] Job stop requested, but cluster stop request failed "
+            f"(status {cluster_response.status_code})."
+        )
+
+
 def _publish_job_asset(
     asset_type: str,
     endpoint_collection: str,
@@ -832,6 +865,7 @@ def command_job_stop(
         response = api.get(f"/experiment/{current_experiment}/jobs/{job_id}/stop")
 
     if response.status_code == 200:
+        _stop_provider_cluster(current_experiment, job_id)
         console.print(f"[success]✓[/success] Job [bold]{job_id}[/bold] stopped.")
     else:
         console.print(f"[error]Error:[/error] Failed to stop job. Status code: {response.status_code}")

cli/tests/commands/test_job.py

@@ -287,6 +287,60 @@ def test_job_info_pretty_still_works(_mock_require, _mock_api):
     assert "RUNNING" in out
 
 
+# ---------------------------------------------------------------------------
+# Stop command tests
+# ---------------------------------------------------------------------------
+
+
+@patch("transformerlab_cli.commands.job.api.post")
+@patch("transformerlab_cli.commands.job.require_current_experiment", return_value="exp1")
+@patch("transformerlab_cli.commands.job.api.get")
+def test_job_stop_also_sends_cluster_stop_when_provider_metadata_exists(_mock_get, _mock_require, mock_post):
+    """job stop should mirror GUI behavior and request provider cluster stop when possible."""
+    stop_response = MagicMock()
+    stop_response.status_code = 200
+    stop_response.json.return_value = {"message": "OK"}
+
+    job_response = MagicMock()
+    job_response.status_code = 200
+    job_response.json.return_value = {
+        "id": "42",
+        "job_data": {"provider_id": "provider-1", "cluster_name": "cluster-a"},
+    }
+
+    _mock_get.side_effect = [stop_response, job_response]
+
+    cluster_stop_response = MagicMock()
+    cluster_stop_response.status_code = 200
+    mock_post.return_value = cluster_stop_response
+
+    result = runner.invoke(app, ["job", "stop", "42"])
+    assert result.exit_code == 0, result.output
+    assert _mock_get.call_args_list[0].args[0] == "/experiment/exp1/jobs/42/stop"
+    assert _mock_get.call_args_list[1].args[0] == "/experiment/exp1/jobs/42"
+    mock_post.assert_called_once_with("/compute_provider/providers/provider-1/clusters/cluster-a/stop?job_id=42")
+
+
+@patch("transformerlab_cli.commands.job.api.post")
+@patch("transformerlab_cli.commands.job.require_current_experiment", return_value="exp1")
+@patch("transformerlab_cli.commands.job.api.get")
+def test_job_stop_skips_cluster_stop_without_provider_metadata(_mock_get, _mock_require, mock_post):
+    """job stop should not call cluster-stop when provider metadata is missing."""
+    stop_response = MagicMock()
+    stop_response.status_code = 200
+    stop_response.json.return_value = {"message": "OK"}
+
+    job_response = MagicMock()
+    job_response.status_code = 200
+    job_response.json.return_value = {"id": "42", "job_data": {}}
+
+    _mock_get.side_effect = [stop_response, job_response]
+
+    result = runner.invoke(app, ["job", "stop", "42"])
+    assert result.exit_code == 0, result.output
+    mock_post.assert_not_called()
+
+
 # ---------------------------------------------------------------------------
 # Log command tests
 # ---------------------------------------------------------------------------