From 6e41890f8dad893257b365012c48bdf5b74c33f6 Mon Sep 17 00:00:00 2001 From: John Thorpe Date: Mon, 31 Aug 2026 10:59:11 -0700 Subject: [PATCH 1/3] Add --availability-model to truss train update A queued training job's capacity guarantee can now be changed in place: truss train update --job-id --availability-model spot Previously a dedicated job blocked on full capacity had to be resubmitted to run on spot. It can now be switched while it sits in the queue, alongside the existing --priority update. Follows the same three layers as the priority update: the option on the `train update` command, the enum-typed parameter in cli/train/core.py, and the PATCH body field in remote/baseten/api.py. Choices are derived from the AvailabilityModel enum so they cannot drift from it, and a choice option rather than a bare --spot flag (as `truss train push` has) since an update also needs to switch a job back to dedicated. Co-Authored-By: Claude Opus 5 --- truss/cli/train/core.py | 13 +++-- truss/cli/train_commands.py | 26 ++++++++-- truss/remote/baseten/api.py | 9 +++- truss/tests/cli/train/test_train_cli_core.py | 52 ++++++++++++++++++++ truss/tests/remote/baseten/test_api.py | 19 +++++++ 5 files changed, 111 insertions(+), 8 deletions(-) diff --git a/truss/cli/train/core.py b/truss/cli/train/core.py index 7d427157c..646c25357 100644 --- a/truss/cli/train/core.py +++ b/truss/cli/train/core.py @@ -184,15 +184,22 @@ def recreate_training_job( def update_training_job( - remote_provider: BasetenRemote, job_id: str, *, priority: Optional[int] = None + remote_provider: BasetenRemote, + job_id: str, + *, + priority: Optional[int] = None, + availability_model: Optional[AvailabilityModel] = None, ) -> Dict[str, Any]: - if priority is None: + if priority is None and availability_model is None: raise ValueError("At least one field to update must be provided.") job = _get_job_by_job_id(remote_provider, job_id) project_id = job["training_project"]["id"] job_id = job["id"] return remote_provider.api.update_training_job( - project_id, job_id, priority=priority + project_id, + job_id, + priority=priority, + availability_model=availability_model.value if availability_model else None, ) diff --git a/truss/cli/train_commands.py b/truss/cli/train_commands.py index 8f64e5009..166026373 100644 --- a/truss/cli/train_commands.py +++ b/truss/cli/train_commands.py @@ -45,7 +45,7 @@ from truss.remote.baseten.remote import BasetenRemote from truss.remote.remote_factory import RemoteFactory from truss.util.path import copy_tree_path -from truss_train import TrainingJob +from truss_train import AvailabilityModel, TrainingJob @click.group() @@ -936,12 +936,25 @@ def update_session( required=False, help="New queue priority. Higher values are dequeued first. Only PENDING jobs can have their priority changed.", ) +@click.option( + "--availability-model", + type=click.Choice([m.value for m in AvailabilityModel], case_sensitive=False), + required=False, + help="New capacity guarantee. 'dedicated' runs on on-demand capacity that is not preempted; " + "'spot' runs on interruptible capacity that may be preempted, and you are responsible for " + "checkpointing your own progress. Only PENDING jobs can have their availability model changed.", +) @click.option("--remote", type=str, required=False, help="Remote to use.") @common.common_options() -def update(job_id: str, priority: Optional[int], remote: Optional[str]): +def update( + job_id: str, + priority: Optional[int], + availability_model: Optional[str], + remote: Optional[str], +): """Update a training job. At least one field to update must be provided.""" - if priority is None: + if priority is None and availability_model is None: raise click.UsageError("At least one field to update must be provided.") if not remote: @@ -953,7 +966,12 @@ def update(job_id: str, priority: Optional[int], remote: Optional[str]): try: job = train_cli.update_training_job( - remote_provider=remote_provider, job_id=job_id, priority=priority + remote_provider=remote_provider, + job_id=job_id, + priority=priority, + availability_model=AvailabilityModel(availability_model.lower()) + if availability_model + else None, ) except Exception as e: error_console.print(f"Failed to update training job: {str(e)}") diff --git a/truss/remote/baseten/api.py b/truss/remote/baseten/api.py index 2cde530fe..f2584e610 100644 --- a/truss/remote/baseten/api.py +++ b/truss/remote/baseten/api.py @@ -912,11 +912,18 @@ def recreate_training_job(self, project_id: str, job_id: str): return resp_json["training_job"] def update_training_job( - self, project_id: str, job_id: str, *, priority: Optional[int] = None + self, + project_id: str, + job_id: str, + *, + priority: Optional[int] = None, + availability_model: Optional[str] = None, ): body: Dict[str, Any] = {} if priority is not None: body["priority"] = priority + if availability_model is not None: + body["availability_model"] = availability_model resp_json = self._rest_api_client.patch( f"v1/training_projects/{project_id}/jobs/{job_id}", body=body ) diff --git a/truss/tests/cli/train/test_train_cli_core.py b/truss/tests/cli/train/test_train_cli_core.py index 4dbaf4723..999b6cc88 100644 --- a/truss/tests/cli/train/test_train_cli_core.py +++ b/truss/tests/cli/train/test_train_cli_core.py @@ -14,9 +14,11 @@ display_training_capacity, display_training_jobs, update_team_training_gpu_capacity, + update_training_job, view_training_job_metrics, ) from truss.remote.baseten.custom_types import FileSummary, TeamType +from truss_train.definitions import AvailabilityModel @patch("truss.cli.train.metrics_watcher.time.sleep") @@ -645,6 +647,56 @@ def test_update_team_training_gpu_capacity_unknown_team_raises(): mock_api.update_team_training_gpu_capacity.assert_not_called() +def _mock_remote_for_job_update(): + mock_api = Mock() + mock_api.search_training_jobs.return_value = [ + {"id": "job_id", "training_project": {"id": "project_id"}} + ] + mock_api.update_training_job.return_value = {"id": "job_id"} + mock_remote = Mock() + mock_remote.api = mock_api + return mock_remote, mock_api + + +def test_update_training_job_availability_model(): + """The availability model is passed to the API as its wire value.""" + mock_remote, mock_api = _mock_remote_for_job_update() + + update_training_job( + mock_remote, "job_id", availability_model=AvailabilityModel.SPOT + ) + + mock_api.update_training_job.assert_called_once_with( + "project_id", "job_id", priority=None, availability_model="spot" + ) + + +def test_update_training_job_priority_and_availability_model(): + """Both fields are forwarded together when both are provided.""" + mock_remote, mock_api = _mock_remote_for_job_update() + + update_training_job( + mock_remote, + "job_id", + priority=7, + availability_model=AvailabilityModel.DEDICATED, + ) + + mock_api.update_training_job.assert_called_once_with( + "project_id", "job_id", priority=7, availability_model="dedicated" + ) + + +def test_update_training_job_no_fields_raises(): + """Updating with no fields raises before any API call is made.""" + mock_remote, mock_api = _mock_remote_for_job_update() + + with pytest.raises(ValueError, match="At least one field"): + update_training_job(mock_remote, "job_id") + + mock_api.update_training_job.assert_not_called() + + def test_format_capacity_type(): """availability_model maps to human-readable labels; absent reads as on-demand.""" assert _format_capacity_type({"availability_model": "spot"}) == "Spot" diff --git a/truss/tests/remote/baseten/test_api.py b/truss/tests/remote/baseten/test_api.py index 33958db7b..8e41a22ef 100644 --- a/truss/tests/remote/baseten/test_api.py +++ b/truss/tests/remote/baseten/test_api.py @@ -598,6 +598,25 @@ def test_update_training_job(mock_patch, baseten_api): assert mock_patch.call_args[1]["json"] == {"priority": 42} +@mock.patch("requests.patch", return_value=mock_update_training_job_response()) +def test_update_training_job_availability_model(mock_patch, baseten_api): + baseten_api.update_training_job("project_id", "job_id", availability_model="spot") + + assert mock_patch.call_args[1]["json"] == {"availability_model": "spot"} + + +@mock.patch("requests.patch", return_value=mock_update_training_job_response()) +def test_update_training_job_priority_and_availability_model(mock_patch, baseten_api): + baseten_api.update_training_job( + "project_id", "job_id", priority=42, availability_model="dedicated" + ) + + assert mock_patch.call_args[1]["json"] == { + "priority": 42, + "availability_model": "dedicated", + } + + # Mock responses for training job logs pagination tests def mock_training_job_logs_response(logs, has_more=True): """Helper function to create mock training job logs response""" From 557fb907b2abdc7be1dc07095d3b0a2db3de7731 Mon Sep 17 00:00:00 2001 From: John Thorpe Date: Mon, 31 Aug 2026 11:15:53 -0700 Subject: [PATCH 2/3] Consolidate update_training_job tests and fix broken call assertion MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit TestUpdateTrainingJob in truss-train/tests/test_recreate.py asserts the exact kwargs passed to the API, so adding availability_model to the call broke test_update_training_job_success. Fix the assertion, and move the two availability-model cases from truss/tests/cli/train/test_train_cli_core.py into TestUpdateTrainingJob beside their siblings — the no-fields case there was an exact duplicate of one already in that class. All five update_training_job call assertions now live in one file, so a signature change can't pass one suite while breaking the other. Co-Authored-By: Claude Opus 5 --- truss-train/tests/test_recreate.py | 44 ++++++++++++++++- truss/tests/cli/train/test_train_cli_core.py | 52 -------------------- 2 files changed, 43 insertions(+), 53 deletions(-) diff --git a/truss-train/tests/test_recreate.py b/truss-train/tests/test_recreate.py index cbf4b9e90..a927219ff 100644 --- a/truss-train/tests/test_recreate.py +++ b/truss-train/tests/test_recreate.py @@ -4,6 +4,7 @@ import pytest from truss.cli.train import core as train_cli +from truss_train.definitions import AvailabilityModel @pytest.fixture @@ -156,7 +157,48 @@ def test_update_training_job_success(self, mock_remote): job_id="test_job_123" ) mock_remote.api.update_training_job.assert_called_once_with( - "project_456", "test_job_123", priority=42 + "project_456", "test_job_123", priority=42, availability_model=None + ) + + def test_update_training_job_availability_model(self, mock_remote): + """The availability model is passed to the API as its wire value.""" + mock_remote.api.search_training_jobs.return_value = [ + { + "id": "test_job_123", + "training_project": {"id": "project_456", "name": "test-project"}, + } + ] + mock_remote.api.update_training_job.return_value = {"id": "test_job_123"} + + train_cli.update_training_job( + remote_provider=mock_remote, + job_id="test_job_123", + availability_model=AvailabilityModel.SPOT, + ) + + mock_remote.api.update_training_job.assert_called_once_with( + "project_456", "test_job_123", priority=None, availability_model="spot" + ) + + def test_update_training_job_priority_and_availability_model(self, mock_remote): + """Both fields are forwarded together when both are provided.""" + mock_remote.api.search_training_jobs.return_value = [ + { + "id": "test_job_123", + "training_project": {"id": "project_456", "name": "test-project"}, + } + ] + mock_remote.api.update_training_job.return_value = {"id": "test_job_123"} + + train_cli.update_training_job( + remote_provider=mock_remote, + job_id="test_job_123", + priority=7, + availability_model=AvailabilityModel.DEDICATED, + ) + + mock_remote.api.update_training_job.assert_called_once_with( + "project_456", "test_job_123", priority=7, availability_model="dedicated" ) def test_update_training_job_no_job_found(self, mock_remote): diff --git a/truss/tests/cli/train/test_train_cli_core.py b/truss/tests/cli/train/test_train_cli_core.py index 999b6cc88..4dbaf4723 100644 --- a/truss/tests/cli/train/test_train_cli_core.py +++ b/truss/tests/cli/train/test_train_cli_core.py @@ -14,11 +14,9 @@ display_training_capacity, display_training_jobs, update_team_training_gpu_capacity, - update_training_job, view_training_job_metrics, ) from truss.remote.baseten.custom_types import FileSummary, TeamType -from truss_train.definitions import AvailabilityModel @patch("truss.cli.train.metrics_watcher.time.sleep") @@ -647,56 +645,6 @@ def test_update_team_training_gpu_capacity_unknown_team_raises(): mock_api.update_team_training_gpu_capacity.assert_not_called() -def _mock_remote_for_job_update(): - mock_api = Mock() - mock_api.search_training_jobs.return_value = [ - {"id": "job_id", "training_project": {"id": "project_id"}} - ] - mock_api.update_training_job.return_value = {"id": "job_id"} - mock_remote = Mock() - mock_remote.api = mock_api - return mock_remote, mock_api - - -def test_update_training_job_availability_model(): - """The availability model is passed to the API as its wire value.""" - mock_remote, mock_api = _mock_remote_for_job_update() - - update_training_job( - mock_remote, "job_id", availability_model=AvailabilityModel.SPOT - ) - - mock_api.update_training_job.assert_called_once_with( - "project_id", "job_id", priority=None, availability_model="spot" - ) - - -def test_update_training_job_priority_and_availability_model(): - """Both fields are forwarded together when both are provided.""" - mock_remote, mock_api = _mock_remote_for_job_update() - - update_training_job( - mock_remote, - "job_id", - priority=7, - availability_model=AvailabilityModel.DEDICATED, - ) - - mock_api.update_training_job.assert_called_once_with( - "project_id", "job_id", priority=7, availability_model="dedicated" - ) - - -def test_update_training_job_no_fields_raises(): - """Updating with no fields raises before any API call is made.""" - mock_remote, mock_api = _mock_remote_for_job_update() - - with pytest.raises(ValueError, match="At least one field"): - update_training_job(mock_remote, "job_id") - - mock_api.update_training_job.assert_not_called() - - def test_format_capacity_type(): """availability_model maps to human-readable labels; absent reads as on-demand.""" assert _format_capacity_type({"availability_model": "spot"}) == "Spot" From 6786dd8261362fd4dfc32e6617b49df19f50d4f9 Mon Sep 17 00:00:00 2001 From: John Thorpe Date: Mon, 31 Aug 2026 14:15:02 -0700 Subject: [PATCH 3/3] Move update_training_job tests out of test_recreate.py test_recreate.py is about job recreation, and it lives in truss-train/tests/ even though it exercises truss.cli.train.core. TestUpdateTrainingJob moves to truss/tests/cli/train/test_job_update.py, which names what it tests and sits in the tree for the package under test. All five update assertions stay together in the new file, so a change to the call signature still can't pass one suite while breaking another. Co-Authored-By: Claude Opus 5 --- truss-train/tests/test_recreate.py | 99 -------------------- truss/tests/cli/train/test_job_update.py | 111 +++++++++++++++++++++++ 2 files changed, 111 insertions(+), 99 deletions(-) create mode 100644 truss/tests/cli/train/test_job_update.py diff --git a/truss-train/tests/test_recreate.py b/truss-train/tests/test_recreate.py index a927219ff..9a68be634 100644 --- a/truss-train/tests/test_recreate.py +++ b/truss-train/tests/test_recreate.py @@ -4,7 +4,6 @@ import pytest from truss.cli.train import core as train_cli -from truss_train.definitions import AvailabilityModel @pytest.fixture @@ -127,101 +126,3 @@ def test_recreate_training_job_without_job_id_user_cancels(self, mock_remote): # Should raise UsageError when user cancels with pytest.raises(click.UsageError, match="Training job not recreated"): train_cli.recreate_training_job(remote_provider=mock_remote) - - -class TestUpdateTrainingJob: - """Test cases for the update_training_job function.""" - - def test_update_training_job_success(self, mock_remote): - """Test updating a training job with a specific job ID.""" - mock_remote.api.search_training_jobs.return_value = [ - { - "id": "test_job_123", - "training_project": {"id": "project_456", "name": "test-project"}, - } - ] - mock_remote.api.update_training_job.return_value = { - "id": "test_job_123", - "priority": 42, - "training_project": {"id": "project_456", "name": "test-project"}, - } - - result = train_cli.update_training_job( - remote_provider=mock_remote, job_id="test_job_123", priority=42 - ) - - assert result["id"] == "test_job_123" - assert result["priority"] == 42 - - mock_remote.api.search_training_jobs.assert_called_once_with( - job_id="test_job_123" - ) - mock_remote.api.update_training_job.assert_called_once_with( - "project_456", "test_job_123", priority=42, availability_model=None - ) - - def test_update_training_job_availability_model(self, mock_remote): - """The availability model is passed to the API as its wire value.""" - mock_remote.api.search_training_jobs.return_value = [ - { - "id": "test_job_123", - "training_project": {"id": "project_456", "name": "test-project"}, - } - ] - mock_remote.api.update_training_job.return_value = {"id": "test_job_123"} - - train_cli.update_training_job( - remote_provider=mock_remote, - job_id="test_job_123", - availability_model=AvailabilityModel.SPOT, - ) - - mock_remote.api.update_training_job.assert_called_once_with( - "project_456", "test_job_123", priority=None, availability_model="spot" - ) - - def test_update_training_job_priority_and_availability_model(self, mock_remote): - """Both fields are forwarded together when both are provided.""" - mock_remote.api.search_training_jobs.return_value = [ - { - "id": "test_job_123", - "training_project": {"id": "project_456", "name": "test-project"}, - } - ] - mock_remote.api.update_training_job.return_value = {"id": "test_job_123"} - - train_cli.update_training_job( - remote_provider=mock_remote, - job_id="test_job_123", - priority=7, - availability_model=AvailabilityModel.DEDICATED, - ) - - mock_remote.api.update_training_job.assert_called_once_with( - "project_456", "test_job_123", priority=7, availability_model="dedicated" - ) - - def test_update_training_job_no_job_found(self, mock_remote): - """Test updating a non-existent job ID.""" - mock_remote.api.search_training_jobs.return_value = [] - - with pytest.raises( - RuntimeError, match="No training job found with ID: nonexistent_job" - ): - train_cli.update_training_job( - remote_provider=mock_remote, job_id="nonexistent_job", priority=42 - ) - - mock_remote.api.update_training_job.assert_not_called() - - def test_update_training_job_no_fields_raises(self, mock_remote): - """Test that updating with no fields provided raises an error.""" - with pytest.raises( - ValueError, match="At least one field to update must be provided" - ): - train_cli.update_training_job( - remote_provider=mock_remote, job_id="test_job_123" - ) - - mock_remote.api.search_training_jobs.assert_not_called() - mock_remote.api.update_training_job.assert_not_called() diff --git a/truss/tests/cli/train/test_job_update.py b/truss/tests/cli/train/test_job_update.py new file mode 100644 index 000000000..fa8faab58 --- /dev/null +++ b/truss/tests/cli/train/test_job_update.py @@ -0,0 +1,111 @@ +from unittest.mock import MagicMock + +import pytest + +from truss.cli.train import core as train_cli +from truss_train.definitions import AvailabilityModel + + +@pytest.fixture +def mock_remote(): + remote = MagicMock() + remote.api = MagicMock() + return remote + + +class TestUpdateTrainingJob: + """Test cases for the update_training_job function.""" + + def test_update_training_job_success(self, mock_remote): + """Test updating a training job with a specific job ID.""" + mock_remote.api.search_training_jobs.return_value = [ + { + "id": "test_job_123", + "training_project": {"id": "project_456", "name": "test-project"}, + } + ] + mock_remote.api.update_training_job.return_value = { + "id": "test_job_123", + "priority": 42, + "training_project": {"id": "project_456", "name": "test-project"}, + } + + result = train_cli.update_training_job( + remote_provider=mock_remote, job_id="test_job_123", priority=42 + ) + + assert result["id"] == "test_job_123" + assert result["priority"] == 42 + + mock_remote.api.search_training_jobs.assert_called_once_with( + job_id="test_job_123" + ) + mock_remote.api.update_training_job.assert_called_once_with( + "project_456", "test_job_123", priority=42, availability_model=None + ) + + def test_update_training_job_availability_model(self, mock_remote): + """The availability model is passed to the API as its wire value.""" + mock_remote.api.search_training_jobs.return_value = [ + { + "id": "test_job_123", + "training_project": {"id": "project_456", "name": "test-project"}, + } + ] + mock_remote.api.update_training_job.return_value = {"id": "test_job_123"} + + train_cli.update_training_job( + remote_provider=mock_remote, + job_id="test_job_123", + availability_model=AvailabilityModel.SPOT, + ) + + mock_remote.api.update_training_job.assert_called_once_with( + "project_456", "test_job_123", priority=None, availability_model="spot" + ) + + def test_update_training_job_priority_and_availability_model(self, mock_remote): + """Both fields are forwarded together when both are provided.""" + mock_remote.api.search_training_jobs.return_value = [ + { + "id": "test_job_123", + "training_project": {"id": "project_456", "name": "test-project"}, + } + ] + mock_remote.api.update_training_job.return_value = {"id": "test_job_123"} + + train_cli.update_training_job( + remote_provider=mock_remote, + job_id="test_job_123", + priority=7, + availability_model=AvailabilityModel.DEDICATED, + ) + + mock_remote.api.update_training_job.assert_called_once_with( + "project_456", "test_job_123", priority=7, availability_model="dedicated" + ) + + def test_update_training_job_no_job_found(self, mock_remote): + """Test updating a non-existent job ID.""" + mock_remote.api.search_training_jobs.return_value = [] + + with pytest.raises( + RuntimeError, match="No training job found with ID: nonexistent_job" + ): + train_cli.update_training_job( + remote_provider=mock_remote, job_id="nonexistent_job", priority=42 + ) + + mock_remote.api.update_training_job.assert_not_called() + + def test_update_training_job_no_fields_raises(self, mock_remote): + """Test that updating with no fields provided raises an error.""" + with pytest.raises( + ValueError, match="At least one field to update must be provided" + ): + train_cli.update_training_job( + remote_provider=mock_remote, job_id="test_job_123" + ) + + mock_remote.api.search_training_jobs.assert_not_called() + mock_remote.api.update_training_job.assert_not_called()