truss

repository·main·Indexed 22 days ago

https://github.com/basetenlabs/truss

A CLI tool and library designed to simplify the deployment and serving of AI/ML models on Baseten. It handles containerization, dependency management, and GPU configuration. The package also includes the @basetenlabs/performance-client for Node.js, providing high-performance capabilities for embeddings, reranking, and classification via a specialized PerformanceClient with support for HTTP/2, request hedging, and custom proxy configurations.

Tokens
75.9K
Snippets
254
Records
339
Agent score
74%

What's inside truss

  1. Overview of the Baseten Performance Client

    main

    The baseten_performance_client is a high-performance Python library designed for massive concurrent POST requests to Baseten.co endpoints (embeddings, reranking, classification) or any other URL.

    Key features include:

    • High Throughput: Benchmarked at >1200 requests per second (rps) per client.
    • GIL Release: The client releases the Python Global Interpreter Lock (GIL) while performing requests in Rust, allowing for true parallelism.
    • Hybrid Usage: Supports simultaneous synchronous and asynchronous usage.
    • Core Stack: Built on pyo3, reqwest, and tokio.
  2. Implement a non-streaming Model class

    main

    A standard Truss model requires a Model class in model/model.py with three specific methods:

    1. __init__(self, **kwargs): Initializes the class instance.
    2. load(self): Runs once when the server starts to load weights, tokenizers, or pipelines into memory.
    3. predict(self, request: Dict) -> Dict: Handles inference. For non-streaming models, this method should accept a dictionary and return a JSON-serializable dictionary.
    class Model:
        def __init__(self, **kwargs) -> None:
            self.tokenizer = None
            self.model = None
    
        def load(self):
            # Load model and tokenizer here
            pass
    
        def predict(self, request: Dict) -> Dict:
            # Perform inference and return a dictionary
            return {"output": "result"}
  3. Implement the Model class in Truss

    main

    Every Truss requires a Model class in model/model.py with three specific member functions:

    • __init__(self, **kwargs): Initializes the object. It must capture kwargs["secrets"] if you need to access Baseten secrets.
    • load(self): Runs once when the model server starts. Use this to load your model (e.g., weights, pipelines) into a class property like self._model.
    • predict(self, model_input): Runs on every inference request. It handles the input and returns a JSON-serializable output.
    from transformers import pipeline
    
    
    class Model:
        def __init__(self, **kwargs) -> None:
            self._secrets = kwargs["secrets"]
            self._model = None
    
        def load(self):
            self._model = pipeline(
                "fill-mask",
                model="baseten/docs-example-gated-model"
            )
    
        def predict(self, model_input):
            return self._model(model_input)
  4. Use Request Hedging for lower latency

    main

    Request hedging allows the client to send duplicate requests after a specified hedge_delay to mitigate tail latency. This is configured via RequestProcessingPreference.

    from baseten_performance_client import RequestProcessingPreference
    
    preference = RequestProcessingPreference(
        hedge_delay=0.5,  # Send hedge request after 0.5s
        max_chars_per_request=256000,
        total_timeout_s=360
    )
    
    response = client.embed(input=texts, model="my_model", preference=preference)
  5. Achieve fast cold starts with weight caching

    main

    To significantly reduce cold start times, you can cache model weights at build time. This bakes the weights directly into the Truss image, making them available immediately when a model replica starts.

    When using this method:

    1. Build time increases: The build process takes longer because it must bundle the model weights.
    2. Deployment and scaling are faster: Once built, the deployment and scale-up processes are much faster because weights do not need to be downloaded at runtime.
  6. How pre/post-process methods work in Truss

    main

    Truss separates IO-bound tasks from compute-bound tasks using preprocess and postprocess methods. This architecture prevents IO operations (like downloading files from a URL) from blocking the compute resources (like a GPU) and allows for higher throughput.

    Concurrency Behavior

    • predict method: Subject to the predict_concurrency limit defined in config.yaml. This ensures the GPU/CPU is not overloaded.
    • preprocess and postprocess methods: These run on separate threads and are not subject to the predict_concurrency limit.

    Example Scenario: If predict_concurrency is set to 5 and you receive 10 requests, all 10 requests will begin preprocess in parallel. However, only 5 will be allowed to enter the predict stage at a time. This ensures that while the model is waiting for IO, the compute engine remains fully utilized.