truss 0.11.10rc5__tar.gz → 0.11.11__tar.gz

{truss-0.11.10rc5 → truss-0.11.11}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: truss
-Version: 0.11.10rc5
+Version: 0.11.11
 Summary: A seamless bridge from model development to model delivery
 Project-URL: Repository, https://github.com/basetenlabs/truss
 Project-URL: Homepage, https://truss.baseten.co

{truss-0.11.10rc5 → truss-0.11.11}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "truss"
-version = "0.11.10rc5"
+version = "0.11.11"
 description = "A seamless bridge from model development to model delivery"
 authors = [
     { name = "Pankaj Gupta", email = "no-reply@baseten.co" },

{truss-0.11.10rc5 → truss-0.11.11}/truss/cli/train/core.py

@@ -274,7 +274,7 @@ def _get_checkpoint_names(
     checkpoint_deploy_config: DeployCheckpointsConfigComplete,
 ) -> list[str]:
     return [
-        checkpoint.paths[0].strip("/").split("/")[-1]
+        checkpoint.checkpoint_name
         for checkpoint in checkpoint_deploy_config.checkpoint_details.checkpoints
     ]
 
@@ -311,9 +311,10 @@ def display_training_job(
     table.add_column("Value")
 
     # Basic job details
+    table.add_row("Job Name", job["name"])
+    table.add_row("Job ID", job["id"])
     table.add_row("Project ID", job["training_project"]["id"])
     table.add_row("Project Name", job["training_project"]["name"])
-    table.add_row("Job ID", job["id"])
     table.add_row("Status", job["current_status"])
     table.add_row("Instance Type", job["instance_type"]["name"])
     table.add_row("Created", cli_common.format_localized_time(job["created_at"]))

{truss-0.11.10rc5 → truss-0.11.11}/truss/cli/train/deploy_checkpoints/deploy_checkpoints.py

@@ -115,18 +115,12 @@ def _build_inference_template_request(
     weights_sources = []
     for checkpoint in checkpoint_deploy_config.checkpoint_details.checkpoints:
         # Extract checkpoint name from the first path
-        checkpoint_name = (
-            checkpoint.paths[0].strip("/").split("/")[-1]
-            if checkpoint.paths
-            else "checkpoint"
-        )
-
         weights_source = {
             "weight_source_type": "B10_CHECKPOINTING",
             "b10_training_checkpoint_weights_source": {
                 "checkpoint": {
                     "training_job_id": checkpoint.training_job_id,
-                    "checkpoint_name": checkpoint_name,
+                    "checkpoint_name": checkpoint.checkpoint_name,
                 }
             },
         }
@@ -309,6 +303,7 @@ def _ensure_checkpoint_details(
     job_id: Optional[str],
 ) -> CheckpointList:
     if checkpoint_details and checkpoint_details.checkpoints:
+        # TODO: check here
         return _process_user_provided_checkpoints(checkpoint_details, remote_provider)
     else:
         return _prompt_user_for_checkpoint_details(

truss-0.11.11/truss/cli/train/deploy_checkpoints/deploy_full_checkpoints.py

@@ -0,0 +1,9 @@
+from truss_train.definitions import FullCheckpoint
+
+
+def hydrate_full_checkpoint(
+    job_id: str, checkpoint_id: str, checkpoint: dict
+) -> FullCheckpoint:
+    """Create a Checkpoint object for full model weights."""
+    # NOTE: Slash at the end is important since it means the checkpoint is a directory
+    return FullCheckpoint(training_job_id=job_id, checkpoint_name=checkpoint_id)

{truss-0.11.10rc5 → truss-0.11.11}/truss/cli/train/deploy_checkpoints/deploy_lora_checkpoints.py

@@ -11,11 +11,10 @@ def hydrate_lora_checkpoint(
 ) -> LoRACheckpoint:
     """Create a LoRA-specific Checkpoint object."""
     # NOTE: Slash at the end is important since it means the checkpoint is a directory
-    paths = [f"rank-0/{checkpoint_id}/"]
     return LoRACheckpoint(
         training_job_id=job_id,
-        paths=paths,
         lora_details=LoRADetails(rank=_get_lora_rank(checkpoint)),
+        checkpoint_name=checkpoint_id,
     )
 
 

{truss-0.11.10rc5 → truss-0.11.11}/truss/cli/train/deploy_checkpoints/deploy_whisper_checkpoints.py

@@ -5,6 +5,4 @@ def hydrate_whisper_checkpoint(
     job_id: str, checkpoint_id: str, checkpoint: dict
 ) -> WhisperCheckpoint:
     """Create a Checkpoint object for whisper model weights."""
-    # NOTE: Slash at the end is important since it means the checkpoint is a directory
-    paths = [f"rank-0/{checkpoint_id}/"]
-    return WhisperCheckpoint(training_job_id=job_id, paths=paths)
+    return WhisperCheckpoint(training_job_id=job_id, checkpoint_name=checkpoint_id)

{truss-0.11.10rc5 → truss-0.11.11}/truss/remote/baseten/api.py

@@ -804,7 +804,6 @@ class BasetenApi:
             }
         }
         """
-
         resp = self._post_graphql_query(
             query_string, variables={"request": request_data}
         )

{truss-0.11.10rc5 → truss-0.11.11}/truss/remote/baseten/service.py

@@ -137,13 +137,6 @@ class BasetenService(TrussService):
 
             return decode_content()
 
-        parsed_response = response.json()
-
-        if "error" in parsed_response:
-            # In the case that the model is in a non-ready state, the response
-            # will be a json with an `error` key.
-            return parsed_response
-
         return response.json()
 
     def authenticate(self) -> dict:

{truss-0.11.10rc5 → truss-0.11.11}/truss/templates/control/control/server.py

@@ -72,9 +72,9 @@ class ControlServer:
             # httptools installed, which does not work with our requests & version
             # of uvicorn.
             http="h11",
+            loop="uvloop",
             **extra_kwargs,
         )
-        cfg.setup_event_loop()
 
         server = uvicorn.Server(cfg)
         asyncio.run(server.serve())

{truss-0.11.10rc5 → truss-0.11.11}/truss/templates/control/requirements.txt

@@ -7,7 +7,6 @@ python-json-logger>=2.0.2
 tenacity>=8.1.0
  # To avoid divergence, this should follow the latest release.
 truss==0.11.1
-# NB(nikhil): Uvicorn 0.36.0 has breaking changes for the event loop, so we pin to a lower version.
-uvicorn>=0.24.0,<0.36.0
+uvicorn>=0.24.0
 uvloop>=0.19.0
 websockets>=10.0

{truss-0.11.10rc5 → truss-0.11.11}/truss/templates/server/requirements.txt

@@ -19,7 +19,6 @@ python-json-logger>=2.0.2
 pyyaml>=6.0.0
 requests>=2.31.0
 truss-transfer==0.0.32
-# NB(nikhil): Uvicorn 0.36.0 has breaking changes for the event loop, so we pin to a lower version.
-uvicorn>=0.24.0,<0.36.0
+uvicorn>=0.24.0
 uvloop>=0.19.0
 websockets>=10.0

{truss-0.11.10rc5 → truss-0.11.11}/truss/templates/server/truss_server.py

@@ -497,9 +497,9 @@ class TrussServer:
             timeout_graceful_shutdown=TIMEOUT_GRACEFUL_SHUTDOWN,
             log_config=log_config.make_log_config(log_level),
             ws_max_size=WS_MAX_MSG_SZ_BYTES,
+            loop="uvloop",
             **extra_kwargs,
         )
-        cfg.setup_event_loop()  # Call this so uvloop gets used
         server = uvicorn.Server(config=cfg)
         self._server = server
         asyncio.run(server.serve())

{truss-0.11.10rc5 → truss-0.11.11}/truss/templates/server.Dockerfile.jinja

@@ -69,7 +69,7 @@ COPY --chown={{ default_owner }} ./{{ config.data_dir }} ${APP_HOME}/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.11rc1/truss-transfer-cli-v0.10.11rc1-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.11.10rc5/truss-transfer-cli-v0.11.10rc5-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-0.11.10rc5 → truss-0.11.11}/truss/tests/cli/train/test_deploy_checkpoints.py

@@ -115,8 +115,7 @@ def test_hydrate_full_checkpoint():
     assert isinstance(result, definitions.FullCheckpoint)
     assert result.training_job_id == job_id
     assert result.model_weight_format == ModelWeightsFormat.FULL
-    assert len(result.paths) == 1
-    assert result.paths[0] == f"rank-0/{checkpoint_id}/"
+    assert result.checkpoint_name == checkpoint_id
 
 
 def test_hydrate_checkpoint_dispatcher_full():
@@ -272,6 +271,6 @@ def test_hydrate_whisper_checkpoint():
     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.checkpoint_name == checkpoint_id
     assert result.model_weight_format == definitions.ModelWeightsFormat.WHISPER
     assert isinstance(result, definitions.WhisperCheckpoint)

{truss-0.11.10rc5 → truss-0.11.11}/truss/tests/remote/baseten/test_service.py

@@ -1,4 +1,7 @@
+from unittest.mock import MagicMock
+
 from truss.remote.baseten import service
+from truss.remote.baseten.core import ModelVersionHandle
 
 
 def test_model_invoke_url_prod():
@@ -65,3 +68,56 @@ def test_chain_logs_url():
         "https://app.baseten.co", "abc", "666", "543"
     )
     assert url == "https://app.baseten.co/chains/abc/logs/666/543"
+
+
+def test_predict_response_to_json():
+    """Test that predict method returns JSON response for normal dict result."""
+    # Create a mock BasetenService
+    mock_handle = MagicMock(spec=ModelVersionHandle)
+    mock_handle.model_id = "test-model"
+    mock_handle.version_id = "test-version"
+    mock_handle.hostname = "https://model-test.api.baseten.co"
+
+    mock_api = MagicMock()
+    mock_api.app_url = "https://app.baseten.co"
+
+    service_instance = service.BasetenService(
+        model_version_handle=mock_handle,
+        is_draft=False,
+        api_key="test-key",
+        service_url="https://test.com",
+        api=mock_api,
+    )
+
+    # Mock the _send_request method to return a successful JSON response
+    mock_response = MagicMock()
+    mock_response.json.return_value = {"result": "success"}
+    service_instance._send_request = MagicMock(return_value=mock_response)
+
+    # Test predict method
+    result = service_instance.predict({"input": "test"})
+
+    # Verify that the JSON response is returned directly
+    assert result == {"result": "success"}
+
+    # Test non-dict response types below
+
+    # With integer response
+    mock_response.json.return_value = 42
+    result = service_instance.predict({"input": "test"})
+    assert result == 42
+
+    # With string response
+    mock_response.json.return_value = "success"
+    result = service_instance.predict({"input": "test"})
+    assert result == "success"
+
+    # With list response
+    mock_response.json.return_value = [1, 2, 3, 4]
+    result = service_instance.predict({"input": "test"})
+    assert result == [1, 2, 3, 4]
+
+    # With boolean response
+    mock_response.json.return_value = True
+    result = service_instance.predict({"input": "test"})
+    assert result is True

{truss-0.11.10rc5 → truss-0.11.11}/truss-train/truss_train/definitions.py

@@ -185,12 +185,13 @@ class TrainingProject(custom_types.SafeModelNoExtra):
 
 class Checkpoint(custom_types.ConfigModel, ABC):
     training_job_id: str
-    paths: List[str]
+    checkpoint_name: str
     model_weight_format: ModelWeightsFormat
 
     def to_truss_config(self) -> truss_config.TrainingArtifactReference:
         return truss_config.TrainingArtifactReference(
-            training_job_id=self.training_job_id, paths=self.paths
+            training_job_id=self.training_job_id,
+            paths=[f"rank-0/{self.checkpoint_name}/"],
         )
 
 

{truss-0.11.10rc5 → truss-0.11.11}/truss-transfer/README.md

@@ -284,7 +284,7 @@ The following environment variables can be used to configure truss-transfer beha
   - Number of concurrent download workers
   - Controls parallelism for file downloads
 
-- **`TRUSS_TRANSFER_USE_RANGE_DOWNLOAD`** (default: `false`)
+- **`TRUSS_TRANSFER_USE_RANGE_DOWNLOAD`** (default: `true`)
   - Enable/disable range-based downloading for large files
   - Set to `1`, `true`, `yes`, or `y` to enable
 

truss-0.11.10rc5/truss/cli/train/deploy_checkpoints/deploy_full_checkpoints.py

@@ -1,35 +0,0 @@
-from pathlib import Path
-
-from truss_train.definitions import FullCheckpoint
-
-
-def hydrate_full_checkpoint(
-    job_id: str, checkpoint_id: str, checkpoint: dict
-) -> FullCheckpoint:
-    """Create a Checkpoint object for full model weights."""
-    # NOTE: Slash at the end is important since it means the checkpoint is a directory
-    paths = [f"rank-0/{checkpoint_id}/"]
-    return FullCheckpoint(training_job_id=job_id, paths=paths)
-
-
-def build_full_checkpoint_string(truss_deploy_config) -> str:
-    """Build checkpoint string from artifact references for full checkpoints.
-
-    Args:
-        truss_deploy_config: The truss deploy configuration containing training checkpoints.
-
-    Returns:
-        A space-separated string of checkpoint paths.
-    """
-    checkpoint_parts = []
-    for (
-        truss_checkpoint
-    ) in truss_deploy_config.training_checkpoints.artifact_references:  # type: ignore
-        ckpt_path = Path(
-            truss_deploy_config.training_checkpoints.download_folder,  # type: ignore
-            truss_checkpoint.training_job_id,
-            truss_checkpoint.paths[0],
-        )
-        checkpoint_parts.append(str(ckpt_path))
-
-    return " ".join(checkpoint_parts)