truss 0.10.9rc601__py3-none-any.whl → 0.10.10__py3-none-any.whl

truss/templates/server/model_wrapper.py

@@ -59,7 +59,6 @@ class MethodName(str, enum.Enum):
     CHAT_COMPLETIONS = enum.auto()
     COMPLETIONS = enum.auto()
     IS_HEALTHY = enum.auto()
-    MESSAGES = enum.auto()
     POSTPROCESS = enum.auto()
     PREDICT = enum.auto()
     PREPROCESS = enum.auto()
@@ -245,7 +244,6 @@ class ModelDescriptor:
     is_healthy: Optional[MethodDescriptor]
     completions: Optional[MethodDescriptor]
     chat_completions: Optional[MethodDescriptor]
-    messages: Optional[MethodDescriptor]
     websocket: Optional[MethodDescriptor]
 
     @cached_property
@@ -293,7 +291,6 @@ class ModelDescriptor:
         setup = cls._safe_extract_descriptor(model_cls, MethodName.SETUP_ENVIRONMENT)
         completions = cls._safe_extract_descriptor(model_cls, MethodName.COMPLETIONS)
         chats = cls._safe_extract_descriptor(model_cls, MethodName.CHAT_COMPLETIONS)
-        messages = cls._safe_extract_descriptor(model_cls, MethodName.MESSAGES)
         is_healthy = cls._safe_extract_descriptor(model_cls, MethodName.IS_HEALTHY)
         if is_healthy and is_healthy.arg_config != ArgConfig.NONE:
             raise errors.ModelDefinitionError(
@@ -362,7 +359,6 @@ class ModelDescriptor:
             is_healthy=is_healthy,
             completions=completions,
             chat_completions=chats,
-            messages=messages,
             websocket=websocket,
         )
 
@@ -929,14 +925,6 @@ class ModelWrapper:
         )
         return await self._execute_model_endpoint(inputs, request, descriptor)
 
-    async def messages(
-        self, inputs: InputType, request: starlette.requests.Request
-    ) -> OutputType:
-        descriptor = self._get_descriptor_or_raise(
-            self.model_descriptor.messages, MethodName.MESSAGES
-        )
-        return await self._execute_model_endpoint(inputs, request, descriptor)
-
     async def websocket(self, ws: WebSocket) -> None:
         descriptor = self.model_descriptor.websocket
         assert descriptor, "websocket can only be invoked if present on model."

truss/templates/server/requirements.txt

@@ -18,7 +18,7 @@ psutil>=5.9.4
 python-json-logger>=2.0.2
 pyyaml>=6.0.0
 requests>=2.31.0
-truss-transfer==0.0.27
+truss-transfer==0.0.29
 uvicorn>=0.24.0
 uvloop>=0.19.0
 websockets>=10.0

truss/templates/server/truss_server.py

@@ -231,13 +231,6 @@ class BasetenEndpoints:
             method=self._model.completions, request=request, body_raw=body_raw
         )
 
-    async def messages(
-        self, request: Request, body_raw: bytes = Depends(parse_body)
-    ) -> Response:
-        return await self._execute_request(
-            method=self._model.messages, request=request, body_raw=body_raw
-        )
-
     async def websocket(self, ws: WebSocket) -> None:
         self.check_healthy()
         trace_ctx = otel_propagate.extract(ws.headers) or None
@@ -435,12 +428,6 @@ class TrussServer:
                     methods=["POST"],
                     tags=["V1"],
                 ),
-                FastAPIRoute(
-                    r"/v1/messages",
-                    self._endpoints.messages,
-                    methods=["POST"],
-                    tags=["V1"],
-                ),
                 # Websocket endpoint
                 FastAPIWebSocketRoute(r"/v1/websocket", self._endpoints.websocket),
                 # Endpoint aliases for Sagemaker hosting

truss/templates/server.Dockerfile.jinja

@@ -70,7 +70,7 @@ COPY ./{{ config.data_dir }} /app/data
 
 {%- if model_cache_v2 %}
 # v0.0.9, keep synced with server_requirements.txt
-RUN curl -sSL --fail --retry 5 --retry-delay 2 -o /usr/local/bin/truss-transfer-cli https://github.com/basetenlabs/truss/releases/download/v0.10.9rc0/truss-transfer-cli-v0.10.9rc0-linux-x86_64-unknown-linux-musl
+RUN curl -sSL --fail --retry 5 --retry-delay 2 -o /usr/local/bin/truss-transfer-cli https://github.com/basetenlabs/truss/releases/download/v0.10.10rc1/truss-transfer-cli-v0.10.10rc1-linux-x86_64-unknown-linux-musl
 RUN chmod +x /usr/local/bin/truss-transfer-cli
 RUN mkdir /static-bptr
 RUN echo "hash {{model_cache_hash}}"

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,
@@ -682,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/contexts/image_builder/test_serving_image_builder.py

@@ -10,7 +10,6 @@ import pytest
 import yaml
 
 from truss.base.constants import (
-    HF_ACCESS_TOKEN_FILE_NAME,
     TRTLLM_BASE_IMAGE,
     TRTLLM_PREDICT_CONCURRENCY,
     TRTLLM_PYTHON_EXECUTABLE,
@@ -18,6 +17,7 @@ from truss.base.constants import (
 )
 from truss.base.truss_config import ModelCache, ModelRepo, TrussConfig
 from truss.contexts.image_builder.serving_image_builder import (
+    HF_ACCESS_TOKEN_FILE_NAME,
     ServingImageBuilderContext,
     get_files_to_model_cache_v1,
 )

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)

truss/tests/remote/baseten/test_api.py

@@ -7,19 +7,10 @@ from requests import Response
 
 import truss_train.definitions as train_definitions
 from truss.remote.baseten import custom_types as b10_types
-from truss.remote.baseten.api import BasetenApi
 from truss.remote.baseten.custom_types import ChainletDataAtomic, OracleData
 from truss.remote.baseten.error import ApiError
 
 
-@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
-
-
 def mock_successful_response():
     response = Response()
     response.status_code = 200
@@ -134,11 +125,6 @@ def mock_deploy_chain_deployment_response():
     return response
 
 
-@pytest.fixture
-def baseten_api(mock_auth_service):
-    return BasetenApi("https://app.test.com", mock_auth_service)
-
-
 @mock.patch("requests.post", return_value=mock_successful_response())
 def test_post_graphql_query_success(mock_post, baseten_api):
     response_data = {"data": {"status": "success"}}
@@ -439,3 +425,52 @@ def test_upsert_training_project(mock_post, baseten_api):
     upsert_body = mock_post.call_args[1]["json"]["training_project"]
     assert "job" not in upsert_body
     assert "training-project" == upsert_body["name"]
+
+
+# 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"""
+    response = Response()
+    response.status_code = 200
+    response.json = mock.Mock(return_value={"logs": logs})
+    return response
+
+
+def mock_training_job_logs_empty_response():
+    """Helper function to create mock empty training job logs response"""
+    response = Response()
+    response.status_code = 200
+    response.json = mock.Mock(return_value={"logs": []})
+    return response
+
+
+def mock_training_job_logs_error_response():
+    """Helper function to create mock error response for training job logs"""
+    response = Response()
+    response.status_code = 500
+    response.raise_for_status = mock.Mock(
+        side_effect=requests.exceptions.HTTPError("Server Error")
+    )
+    return response
+
+
+def test_fetch_log_batch(baseten_api):
+    """Test _fetch_log_batch helper method"""
+
+    mock_logs = [
+        {"timestamp": "1640995200000000000", "message": "Log 1"},
+        {"timestamp": "1640995260000000000", "message": "Log 2"},
+    ]
+
+    # Mock the rest_api_client
+    mock_rest_client = mock.Mock()
+    mock_rest_client.post.return_value = {"logs": mock_logs}
+    baseten_api._rest_api_client = mock_rest_client
+
+    query_params = {"limit": 100, "direction": "asc"}
+    result = baseten_api._fetch_log_batch("project-123", "job-456", query_params)
+
+    assert result == mock_logs
+    mock_rest_client.post.assert_called_with(
+        "v1/training_projects/project-123/jobs/job-456/logs", body=query_params
+    )