truss 0.10.9rc535__py3-none-any.whl → 0.10.10rc0__py3-none-any.whl

truss/tests/cli/train/test_deploy_checkpoints.py

@@ -1,6 +1,8 @@
 import os
 import re
+from dataclasses import dataclass
 from pathlib import Path
+from typing import Dict, List, Optional
 from unittest.mock import MagicMock, patch
 
 import pytest
@@ -23,6 +25,11 @@ from truss.cli.train.deploy_checkpoints.deploy_lora_checkpoints import (
     hydrate_lora_checkpoint,
     render_vllm_lora_truss_config,
 )
+from truss.cli.train.deploy_checkpoints.deploy_whisper_checkpoints import (
+    VLLM_WHISPER_START_COMMAND,
+    hydrate_whisper_checkpoint,
+    render_vllm_whisper_truss_config,
+)
 from truss.cli.train.types import (
     DeployCheckpointsConfigComplete,
     PrepareCheckpointResult,
@@ -505,8 +512,16 @@ def test_render_vllm_full_truss_config():
     )
 
     result = render_vllm_full_truss_config(deploy_config)
-
-    expected_vllm_command = 'sh -c "HF_TOKEN=$(cat /secrets/hf_token) vllm serve /tmp/training_checkpoints/job123/rank-0/checkpoint-1 --port 8000 --tensor-parallel-size 2 --dtype bfloat16"'
+    expected_vllm_command = (
+        "sh -c 'HF_TOKEN=$(cat /secrets/hf_token) "
+        'HF_TOKEN="$$(cat /secrets/hf_access_token)" && export HF_TOKEN && '
+        "if [ -f /tmp/training_checkpoints/job123/rank-0/checkpoint-1/chat_template.jinja ]; then   "
+        "vllm serve /tmp/training_checkpoints/job123/rank-0/checkpoint-1 "
+        "--chat-template /tmp/training_checkpoints/job123/rank-0/checkpoint-1/chat_template.jinja   "
+        "--port 8000 --tensor-parallel-size 2 --dtype bfloat16; else   "
+        "vllm serve /tmp/training_checkpoints/job123/rank-0/checkpoint-1 "
+        "--port 8000 --tensor-parallel-size 2 --dtype bfloat16; fi'"
+    )
 
     assert isinstance(result, truss_config.TrussConfig)
     assert result.model_name == "test-full-model"
@@ -674,3 +689,432 @@ def test_get_checkpoint_ids_to_deploy_single_checkpoint():
 
         # Should return the single checkpoint directly
         assert result == ["checkpoint-1"]
+
+
+def test_vllm_whisper_start_command_template():
+    """Test that the VLLM_WHISPER_START_COMMAND template renders correctly."""
+    # Test with all variables
+    result = VLLM_WHISPER_START_COMMAND.render(
+        model_path="/path/to/model",
+        envvars="CUDA_VISIBLE_DEVICES=0",
+        specify_tensor_parallelism=4,
+    )
+
+    expected = (
+        "sh -c 'CUDA_VISIBLE_DEVICES=0 "
+        'HF_TOKEN="$$(cat /secrets/hf_access_token)" && export HF_TOKEN && '
+        "vllm serve /path/to/model --port 8000 --tensor-parallel-size 4'"
+    )
+    assert result == expected
+
+    result = VLLM_WHISPER_START_COMMAND.render(
+        model_path="/path/to/model", envvars=None, specify_tensor_parallelism=1
+    )
+
+    expected = (
+        "sh -c '"
+        'HF_TOKEN="$$(cat /secrets/hf_access_token)" && export HF_TOKEN && '
+        "vllm serve /path/to/model --port 8000 --tensor-parallel-size 1'"
+    )
+    assert result == expected
+
+
+def test_hydrate_whisper_checkpoint():
+    """Test that hydrate_whisper_checkpoint creates correct WhisperCheckpoint object."""
+    job_id = "test-job-123"
+    checkpoint_id = "checkpoint-456"
+    checkpoint = {"some": "data"}
+
+    result = hydrate_whisper_checkpoint(job_id, checkpoint_id, checkpoint)
+
+    assert result.training_job_id == job_id
+    assert result.paths == [f"rank-0/{checkpoint_id}/"]
+    assert result.model_weight_format == definitions.ModelWeightsFormat.WHISPER
+    assert isinstance(result, definitions.WhisperCheckpoint)
+
+
+@patch(
+    "truss.cli.train.deploy_checkpoints.deploy_whisper_checkpoints.setup_base_truss_config"
+)
+@patch(
+    "truss.cli.train.deploy_checkpoints.deploy_whisper_checkpoints.setup_environment_variables_and_secrets"
+)
+@patch(
+    "truss.cli.train.deploy_checkpoints.deploy_whisper_checkpoints.build_full_checkpoint_string"
+)
+def test_render_vllm_whisper_truss_config(
+    mock_build_full_checkpoint_string, mock_setup_env_vars, mock_setup_base_config
+):
+    """Test that render_vllm_whisper_truss_config renders truss config correctly."""
+    # Mock dependencies
+    mock_truss_config = MagicMock()
+    mock_truss_config.environment_variables = {}
+    mock_truss_config.docker_server = MagicMock()
+    mock_setup_base_config.return_value = mock_truss_config
+
+    mock_setup_env_vars.return_value = "HF_TOKEN=$(cat /secrets/hf_access_token)"
+    mock_build_full_checkpoint_string.return_value = "/path/to/checkpoint"
+
+    # Create test config
+    deploy_config = DeployCheckpointsConfigComplete(
+        checkpoint_details=definitions.CheckpointList(
+            checkpoints=[
+                definitions.WhisperCheckpoint(
+                    training_job_id="job123",
+                    paths=["rank-0/checkpoint-1/"],
+                    model_weight_format=definitions.ModelWeightsFormat.WHISPER,
+                )
+            ],
+            base_model_id="openai/whisper-large-v3",
+        ),
+        model_name="whisper-large-v3-vLLM",
+        compute=definitions.Compute(
+            accelerator=truss_config.AcceleratorSpec(accelerator="H100", count=4)
+        ),
+        runtime=definitions.DeployCheckpointsRuntime(
+            environment_variables={
+                "HF_TOKEN": definitions.SecretReference(name="hf_access_token")
+            }
+        ),
+        deployment_name="whisper-large-v3-vLLM",
+        model_weight_format=definitions.ModelWeightsFormat.WHISPER,
+    )
+
+    result = render_vllm_whisper_truss_config(deploy_config)
+
+    mock_setup_base_config.assert_called_once_with(deploy_config)
+    mock_setup_env_vars.assert_called_once_with(mock_truss_config, deploy_config)
+    mock_build_full_checkpoint_string.assert_called_once_with(mock_truss_config)
+
+    assert result == mock_truss_config
+
+    expected_start_command = (
+        "sh -c 'HF_TOKEN=$(cat /secrets/hf_access_token) "
+        'HF_TOKEN="$$(cat /secrets/hf_access_token)" && export HF_TOKEN && '
+        "vllm serve /path/to/checkpoint --port 8000 --tensor-parallel-size 4'"
+    )
+    assert (
+        result.environment_variables[START_COMMAND_ENVVAR_NAME]
+        == expected_start_command
+    )
+
+    assert result.docker_server.start_command == f"%(ENV_{START_COMMAND_ENVVAR_NAME})s"
+
+
+@patch(
+    "truss.cli.train.deploy_checkpoints.deploy_whisper_checkpoints.setup_base_truss_config"
+)
+@patch(
+    "truss.cli.train.deploy_checkpoints.deploy_whisper_checkpoints.setup_environment_variables_and_secrets"
+)
+@patch(
+    "truss.cli.train.deploy_checkpoints.deploy_whisper_checkpoints.build_full_checkpoint_string"
+)
+def test_render_vllm_whisper_truss_config_with_envvars(
+    mock_build_full_checkpoint_string, mock_setup_env_vars, mock_setup_base_config
+):
+    """Test that render_vllm_whisper_truss_config handles environment variables correctly."""
+    # Mock dependencies
+    mock_truss_config = MagicMock()
+    mock_truss_config.environment_variables = {}
+    mock_truss_config.docker_server = MagicMock()
+    mock_setup_base_config.return_value = mock_truss_config
+
+    mock_setup_env_vars.return_value = "CUDA_VISIBLE_DEVICES=0,1"
+    mock_build_full_checkpoint_string.return_value = "/path/to/checkpoint"
+
+    # Create test config with environment variables
+    deploy_config = DeployCheckpointsConfigComplete(
+        checkpoint_details=definitions.CheckpointList(
+            checkpoints=[
+                definitions.WhisperCheckpoint(
+                    training_job_id="job123",
+                    paths=["rank-0/checkpoint-1/"],
+                    model_weight_format=definitions.ModelWeightsFormat.WHISPER,
+                )
+            ],
+            base_model_id="openai/whisper-large-v3",
+        ),
+        model_name="whisper-large-v3-vLLM",
+        compute=definitions.Compute(
+            accelerator=truss_config.AcceleratorSpec(accelerator="H100", count=2)
+        ),
+        runtime=definitions.DeployCheckpointsRuntime(
+            environment_variables={
+                "CUDA_VISIBLE_DEVICES": "0,1",
+                "HF_TOKEN": definitions.SecretReference(name="hf_access_token"),
+            }
+        ),
+        deployment_name="whisper-large-v3-vLLM",
+        model_weight_format=definitions.ModelWeightsFormat.WHISPER,
+    )
+
+    # Call function under test
+    result = render_vllm_whisper_truss_config(deploy_config)
+
+    # Verify environment variables are included in start command
+    expected_start_command = (
+        "sh -c 'CUDA_VISIBLE_DEVICES=0,1 "
+        'HF_TOKEN="$$(cat /secrets/hf_access_token)" && export HF_TOKEN && '
+        "vllm serve /path/to/checkpoint --port 8000 --tensor-parallel-size 2'"
+    )
+    assert (
+        result.environment_variables[START_COMMAND_ENVVAR_NAME]
+        == expected_start_command
+    )
+
+
+@dataclass
+class TestCase:
+    """Test case for setup_base_truss_config function."""
+
+    desc: str
+    input_config: DeployCheckpointsConfigComplete
+    expected_model_name: str
+    expected_predict_endpoint: str
+    expected_accelerator: Optional[str]
+    expected_accelerator_count: Optional[int]
+    expected_checkpoint_paths: List[str]
+    expected_environment_variables: Dict[str, str]
+    should_raise: Optional[str] = None  # Error message if function should raise
+
+    __test__ = False  # Tell pytest this is not a test class
+
+
+def test_setup_base_truss_config():
+    """Table-driven test for setup_base_truss_config function."""
+    from truss.cli.train.deploy_checkpoints.deploy_checkpoints_helpers import (
+        setup_base_truss_config,
+    )
+
+    # Define test cases
+    test_cases = [
+        TestCase(
+            desc="LoRA checkpoint with H100 accelerator",
+            input_config=DeployCheckpointsConfigComplete(
+                checkpoint_details=definitions.CheckpointList(
+                    checkpoints=[
+                        definitions.LoRACheckpoint(
+                            training_job_id="job123",
+                            paths=["rank-0/checkpoint-1/"],
+                            model_weight_format=ModelWeightsFormat.LORA,
+                            lora_details=definitions.LoRADetails(rank=32),
+                        )
+                    ],
+                    base_model_id="google/gemma-3-27b-it",
+                ),
+                model_name="test-lora-model",
+                compute=definitions.Compute(
+                    accelerator=truss_config.AcceleratorSpec(
+                        accelerator="H100", count=4
+                    )
+                ),
+                runtime=definitions.DeployCheckpointsRuntime(environment_variables={}),
+                deployment_name="test-deployment",
+                model_weight_format=ModelWeightsFormat.LORA,
+            ),
+            expected_model_name="test-lora-model",
+            expected_predict_endpoint="/v1/chat/completions",
+            expected_accelerator="H100",
+            expected_accelerator_count=4,
+            expected_checkpoint_paths=["rank-0/checkpoint-1/"],
+            expected_environment_variables={
+                "VLLM_LOGGING_LEVEL": "WARNING",
+                "VLLM_USE_V1": "0",
+                "HF_HUB_ENABLE_HF_TRANSFER": "1",
+            },
+        ),
+        TestCase(
+            desc="Whisper checkpoint with A100 accelerator",
+            input_config=DeployCheckpointsConfigComplete(
+                checkpoint_details=definitions.CheckpointList(
+                    checkpoints=[
+                        definitions.WhisperCheckpoint(
+                            training_job_id="job123",
+                            paths=["rank-0/checkpoint-1/"],
+                            model_weight_format=definitions.ModelWeightsFormat.WHISPER,
+                        )
+                    ],
+                    base_model_id="openai/whisper-large-v3",
+                ),
+                model_name="test-whisper-model",
+                compute=definitions.Compute(
+                    accelerator=truss_config.AcceleratorSpec(
+                        accelerator="A100", count=2
+                    )
+                ),
+                runtime=definitions.DeployCheckpointsRuntime(environment_variables={}),
+                deployment_name="test-whisper-deployment",
+                model_weight_format=definitions.ModelWeightsFormat.WHISPER,
+            ),
+            expected_model_name="test-whisper-model",
+            expected_predict_endpoint="/v1/audio/transcriptions",
+            expected_accelerator="A100",
+            expected_accelerator_count=2,
+            expected_checkpoint_paths=["rank-0/checkpoint-1/"],
+            expected_environment_variables={
+                "VLLM_LOGGING_LEVEL": "WARNING",
+                "VLLM_USE_V1": "0",
+                "HF_HUB_ENABLE_HF_TRANSFER": "1",
+            },
+        ),
+        TestCase(
+            desc="Multiple LoRA checkpoints",
+            input_config=DeployCheckpointsConfigComplete(
+                checkpoint_details=definitions.CheckpointList(
+                    checkpoints=[
+                        definitions.LoRACheckpoint(
+                            training_job_id="job123",
+                            paths=["rank-0/checkpoint-1/"],
+                            model_weight_format=ModelWeightsFormat.LORA,
+                            lora_details=definitions.LoRADetails(rank=16),
+                        ),
+                        definitions.LoRACheckpoint(
+                            training_job_id="job123",
+                            paths=["rank-0/checkpoint-2/"],
+                            model_weight_format=ModelWeightsFormat.LORA,
+                            lora_details=definitions.LoRADetails(rank=32),
+                        ),
+                    ],
+                    base_model_id="google/gemma-3-27b-it",
+                ),
+                model_name="test-multi-checkpoint-model",
+                compute=definitions.Compute(
+                    accelerator=truss_config.AcceleratorSpec(
+                        accelerator="H100", count=4
+                    )
+                ),
+                runtime=definitions.DeployCheckpointsRuntime(environment_variables={}),
+                deployment_name="test-multi-deployment",
+                model_weight_format=ModelWeightsFormat.LORA,
+            ),
+            expected_model_name="test-multi-checkpoint-model",
+            expected_predict_endpoint="/v1/chat/completions",
+            expected_accelerator="H100",
+            expected_accelerator_count=4,
+            expected_checkpoint_paths=["rank-0/checkpoint-1/", "rank-0/checkpoint-2/"],
+            expected_environment_variables={
+                "VLLM_LOGGING_LEVEL": "WARNING",
+                "VLLM_USE_V1": "0",
+                "HF_HUB_ENABLE_HF_TRANSFER": "1",
+            },
+        ),
+        TestCase(
+            desc="No accelerator specified",
+            input_config=DeployCheckpointsConfigComplete(
+                checkpoint_details=definitions.CheckpointList(
+                    checkpoints=[
+                        definitions.LoRACheckpoint(
+                            training_job_id="job123",
+                            paths=["rank-0/checkpoint-1/"],
+                            model_weight_format=ModelWeightsFormat.LORA,
+                            lora_details=definitions.LoRADetails(rank=16),
+                        )
+                    ],
+                    base_model_id="google/gemma-3-27b-it",
+                ),
+                model_name="test-no-accelerator-model",
+                compute=definitions.Compute(),  # No accelerator specified
+                runtime=definitions.DeployCheckpointsRuntime(environment_variables={}),
+                deployment_name="test-no-accelerator-deployment",
+                model_weight_format=ModelWeightsFormat.LORA,
+            ),
+            expected_model_name="test-no-accelerator-model",
+            expected_predict_endpoint="/v1/chat/completions",
+            expected_accelerator=None,
+            expected_accelerator_count=None,
+            expected_checkpoint_paths=["rank-0/checkpoint-1/"],
+            expected_environment_variables={
+                "VLLM_LOGGING_LEVEL": "WARNING",
+                "VLLM_USE_V1": "0",
+                "HF_HUB_ENABLE_HF_TRANSFER": "1",
+            },
+        ),
+    ]
+
+    # Run test cases
+    for test_case in test_cases:
+        print(f"Running test case: {test_case.desc}")
+
+        if test_case.should_raise:
+            # Test error cases
+            with pytest.raises(Exception, match=test_case.should_raise):
+                setup_base_truss_config(test_case.input_config)
+        else:
+            # Test success cases
+            result = setup_base_truss_config(test_case.input_config)
+
+            # Verify basic structure
+            assert isinstance(result, truss_config.TrussConfig), (
+                f"Test case '{test_case.desc}': Result should be TrussConfig"
+            )
+            assert result.model_name == test_case.expected_model_name, (
+                f"Test case '{test_case.desc}': Model name mismatch"
+            )
+
+            # Verify docker server configuration
+            assert result.docker_server is not None, (
+                f"Test case '{test_case.desc}': Docker server should not be None"
+            )
+            assert result.docker_server.start_command == 'sh -c ""', (
+                f"Test case '{test_case.desc}': Start command mismatch"
+            )
+            assert result.docker_server.readiness_endpoint == "/health", (
+                f"Test case '{test_case.desc}': Readiness endpoint mismatch"
+            )
+            assert result.docker_server.liveness_endpoint == "/health", (
+                f"Test case '{test_case.desc}': Liveness endpoint mismatch"
+            )
+            assert (
+                result.docker_server.predict_endpoint
+                == test_case.expected_predict_endpoint
+            ), f"Test case '{test_case.desc}': Predict endpoint mismatch"
+            assert result.docker_server.server_port == 8000, (
+                f"Test case '{test_case.desc}': Server port mismatch"
+            )
+
+            # Verify training checkpoints
+            assert result.training_checkpoints is not None, (
+                f"Test case '{test_case.desc}': Training checkpoints should not be None"
+            )
+            assert len(result.training_checkpoints.artifact_references) == len(
+                test_case.expected_checkpoint_paths
+            ), f"Test case '{test_case.desc}': Number of checkpoint artifacts mismatch"
+
+            for i, expected_path in enumerate(test_case.expected_checkpoint_paths):
+                artifact_ref = result.training_checkpoints.artifact_references[i]
+                assert artifact_ref.paths == [expected_path], (
+                    f"Test case '{test_case.desc}': Checkpoint path {i} mismatch"
+                )
+
+            # Verify resources
+            assert result.resources is not None, (
+                f"Test case '{test_case.desc}': Resources should not be None"
+            )
+
+            if test_case.expected_accelerator:
+                assert result.resources.accelerator is not None, (
+                    f"Test case '{test_case.desc}': Accelerator should not be None"
+                )
+                assert (
+                    result.resources.accelerator.accelerator
+                    == test_case.expected_accelerator
+                ), f"Test case '{test_case.desc}': Accelerator type mismatch"
+                assert (
+                    result.resources.accelerator.count
+                    == test_case.expected_accelerator_count
+                ), f"Test case '{test_case.desc}': Accelerator count mismatch"
+            else:
+                # When no accelerator is specified, it creates an AcceleratorSpec with None values
+                assert result.resources.accelerator is not None, (
+                    f"Test case '{test_case.desc}': Accelerator should exist"
+                )
+                assert result.resources.accelerator.accelerator is None, (
+                    f"Test case '{test_case.desc}': Accelerator type should be None"
+                )
+
+            # Verify environment variables
+            for key, expected_value in test_case.expected_environment_variables.items():
+                assert result.environment_variables[key] == expected_value, (
+                    f"Test case '{test_case.desc}': Environment variable {key} mismatch"
+                )

truss/tests/cli/train/test_train_cli_core.py

@@ -51,6 +51,38 @@ def test_view_training_job_metrics(time_sleep, capfd):
                 "0": [{"timestamp": "", "value": 4321}],
                 "1": [{"timestamp": "", "value": 2222}],
             },
+            "per_node_metrics": [
+                {
+                    "node_id": "node-0",
+                    "metrics": {
+                        "cpu_usage": [{"timestamp": "", "value": 3.2}],
+                        "cpu_memory_usage_bytes": [{"timestamp": "", "value": 1234}],
+                        "gpu_utilization": {
+                            "0": [{"timestamp": "", "value": 0.2}],
+                            "1": [{"timestamp": "", "value": 0.3}],
+                        },
+                        "gpu_memory_usage_bytes": {
+                            "0": [{"timestamp": "", "value": 4321}],
+                            "1": [{"timestamp": "", "value": 2222}],
+                        },
+                    },
+                },
+                {
+                    "node_id": "node-1",
+                    "metrics": {
+                        "cpu_usage": [{"timestamp": "", "value": 2.8}],
+                        "cpu_memory_usage_bytes": [{"timestamp": "", "value": 1000}],
+                        "gpu_utilization": {
+                            "0": [{"timestamp": "", "value": 0.15}],
+                            "1": [{"timestamp": "", "value": 0.25}],
+                        },
+                        "gpu_memory_usage_bytes": {
+                            "0": [{"timestamp": "", "value": 4000}],
+                            "1": [{"timestamp": "", "value": 2000}],
+                        },
+                    },
+                },
+            ],
         },
         {
             "training_job": {
@@ -68,6 +100,38 @@ def test_view_training_job_metrics(time_sleep, capfd):
                 "0": [{"timestamp": "", "value": 4321}],
                 "1": [{"timestamp": "", "value": 2222}],
             },
+            "per_node_metrics": [
+                {
+                    "node_id": "node-0",
+                    "metrics": {
+                        "cpu_usage": [{"timestamp": "", "value": 3.2}],
+                        "cpu_memory_usage_bytes": [{"timestamp": "", "value": 1234}],
+                        "gpu_utilization": {
+                            "0": [{"timestamp": "", "value": 0.2}],
+                            "1": [{"timestamp": "", "value": 0.3}],
+                        },
+                        "gpu_memory_usage_bytes": {
+                            "0": [{"timestamp": "", "value": 4321}],
+                            "1": [{"timestamp": "", "value": 2222}],
+                        },
+                    },
+                },
+                {
+                    "node_id": "node-1",
+                    "metrics": {
+                        "cpu_usage": [{"timestamp": "", "value": 2.8}],
+                        "cpu_memory_usage_bytes": [{"timestamp": "", "value": 1000}],
+                        "gpu_utilization": {
+                            "0": [{"timestamp": "", "value": 0.15}],
+                            "1": [{"timestamp": "", "value": 0.25}],
+                        },
+                        "gpu_memory_usage_bytes": {
+                            "0": [{"timestamp": "", "value": 4000}],
+                            "1": [{"timestamp": "", "value": 2000}],
+                        },
+                    },
+                },
+            ],
         },
         {
             "training_job": {
@@ -85,6 +149,38 @@ def test_view_training_job_metrics(time_sleep, capfd):
                 "0": [{"timestamp": "", "value": 4321}],
                 "1": [{"timestamp": "", "value": 2222}],
             },
+            "per_node_metrics": [
+                {
+                    "node_id": "node-0",
+                    "metrics": {
+                        "cpu_usage": [{"timestamp": "", "value": 3.2}],
+                        "cpu_memory_usage_bytes": [{"timestamp": "", "value": 1234}],
+                        "gpu_utilization": {
+                            "0": [{"timestamp": "", "value": 0.2}],
+                            "1": [{"timestamp": "", "value": 0.3}],
+                        },
+                        "gpu_memory_usage_bytes": {
+                            "0": [{"timestamp": "", "value": 4321}],
+                            "1": [{"timestamp": "", "value": 2222}],
+                        },
+                    },
+                },
+                {
+                    "node_id": "node-1",
+                    "metrics": {
+                        "cpu_usage": [{"timestamp": "", "value": 2.8}],
+                        "cpu_memory_usage_bytes": [{"timestamp": "", "value": 1000}],
+                        "gpu_utilization": {
+                            "0": [{"timestamp": "", "value": 0.15}],
+                            "1": [{"timestamp": "", "value": 0.25}],
+                        },
+                        "gpu_memory_usage_bytes": {
+                            "0": [{"timestamp": "", "value": 4000}],
+                            "1": [{"timestamp": "", "value": 2002}],
+                        },
+                    },
+                },
+            ],
         },
     ]
 

truss/tests/remote/baseten/conftest.py

@@ -0,0 +1,18 @@
+from unittest import mock
+
+import pytest
+
+from truss.remote.baseten.api import BasetenApi
+
+
+@pytest.fixture
+def mock_auth_service():
+    auth_service = mock.Mock()
+    auth_token = mock.Mock(headers=lambda: {"Authorization": "Api-Key token"})
+    auth_service.authenticate.return_value = auth_token
+    return auth_service
+
+
+@pytest.fixture
+def baseten_api(mock_auth_service):
+    return BasetenApi("https://app.test.com", mock_auth_service)