Add --availability-model to truss train update

truss-train/tests/test_recreate.py

@@ -126,60 +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
-        )
-
-    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()

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,
     )
 
 

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)}")

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
         )

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()

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"""