truss 0.10.0rc1__py3-none-any.whl → 0.60.0__py3-none-any.whl

truss/tests/test_data/test_streaming_truss_with_tracing/config.yaml

@@ -0,0 +1,43 @@
+apply_library_patches: true
+base_image: null
+build:
+  arguments: {}
+  model_server: TrussServer
+  secret_to_path_mapping: {}
+build_commands: []
+bundled_packages_dir: packages
+data_dir: data
+description: null
+docker_server: null
+environment_variables:
+  OTEL_TRACING_NDJSON_FILE: /tmp/otel_traces.ndjson
+examples_filename: examples.yaml
+external_data: null
+external_package_dirs: []
+input_type: Any
+live_reload: false
+model_cache: []
+model_class_filename: model.py
+model_class_name: Model
+model_framework: custom
+model_metadata: {}
+model_module_dir: model
+model_name: Test Streaming
+model_type: Model
+python_version: py39
+requirements: []
+requirements_file: null
+resources:
+  accelerator: null
+  cpu: '1'
+  memory: 2Gi
+  use_gpu: false
+runtime:
+  enable_tracing_data: false
+  num_workers: 1
+  predict_concurrency: 1
+  streaming_read_timeout: 60
+secrets: {}
+spec_version: '2.0'
+system_packages: []
+trt_llm: null

truss/tests/test_data/test_streaming_truss_with_tracing/model/model.py

@@ -0,0 +1,65 @@
+import json
+import pathlib
+import time
+from typing import Any, Generator, Sequence
+
+import opentelemetry.sdk.resources as resources
+import opentelemetry.sdk.trace as sdk_trace
+import opentelemetry.sdk.trace.export as trace_export
+from opentelemetry import trace
+from opentelemetry.sdk.resources import Resource
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+
+
+class JSONFileExporter(trace_export.SpanExporter):
+    # Copied form truss/templates/server/common/tracing.py.
+
+    def __init__(self, file_path: pathlib.Path):
+        self._file = file_path.open("a")
+
+    def export(
+        self, spans: Sequence[sdk_trace.ReadableSpan]
+    ) -> trace_export.SpanExportResult:
+        for span in spans:
+            # Get rid of newlines and whitespace.
+            self._file.write(json.dumps(json.loads(span.to_json())))
+            self._file.write("\n")
+        self._file.flush()
+        return trace_export.SpanExportResult.SUCCESS
+
+    def shutdown(self) -> None:
+        self._file.close()
+
+
+trace.set_tracer_provider(
+    TracerProvider(resource=Resource.create({resources.SERVICE_NAME: "UserModel"}))
+)
+tracer = trace.get_tracer(__name__)
+file_exporter = JSONFileExporter(pathlib.Path("/tmp/otel_user_traces.ndjson"))
+span_processor = SimpleSpanProcessor(file_exporter)
+trace.get_tracer_provider().add_span_processor(span_processor)
+
+
+class Model:
+    def __init__(self, **kwargs) -> None:
+        self._data_dir = kwargs["data_dir"]
+        self._config = kwargs["config"]
+        self._secrets = kwargs["secrets"]
+        self._model = None
+
+    @tracer.start_as_current_span("load_model")
+    def load(self):
+        pass
+
+    @tracer.start_as_current_span("predict")
+    def predict(self, model_input: Any) -> Generator[str, None, None]:
+        with tracer.start_as_current_span("start-predict") as span:
+
+            def inner():
+                time.sleep(0.02)
+                for i in range(5):
+                    span.add_event("yield")
+                    yield str(i)
+
+            return inner()

truss/tests/test_data/test_trt_llm_truss/config.yaml

@@ -0,0 +1,15 @@
+model_name: Test
+resources:
+  accelerator: A100
+  use_gpu: True
+trt_llm:
+  build:
+    max_seq_len: 1000
+    max_batch_size: 1
+    max_beam_width: 1
+    base_model: llama
+    checkpoint_repository:
+      repo: TinyLlama/TinyLlama-1.1B-Chat-v1.0
+      source: HF
+environment_variables:
+  OTEL_TRACING_NDJSON_FILE: "/tmp/otel_traces.ndjson"

truss/tests/test_data/test_trt_llm_truss/model/model.py

@@ -0,0 +1,15 @@
+from typing import Any
+
+import fastapi
+
+
+class Model:
+    def __init__(self, trt_llm, **kwargs) -> None:
+        self._data_dir = kwargs["data_dir"]
+        self._config = kwargs["config"]
+        self._secrets = kwargs["secrets"]
+        self._model = None
+        self._engine = trt_llm["engine"]
+
+    async def predict(self, model_input: Any, request: fastapi.Request) -> Any:
+        return await self._engine.predict(model_input, request)

truss/tests/test_data/test_truss/config.yaml

@@ -0,0 +1,4 @@
+model_name: Test
+python_version: py39
+environment_variables:
+  OTEL_TRACING_NDJSON_FILE: "/tmp/otel_traces.ndjson"

truss/tests/test_data/test_truss_with_error/config.yaml

@@ -0,0 +1,4 @@
+model_name: Test
+python_version: py39
+environment_variables:
+  OTEL_TRACING_NDJSON_FILE: "/tmp/otel_traces.ndjson"

truss/tests/test_data/test_truss_with_error/model/model.py

@@ -0,0 +1,8 @@
+from typing import Any
+
+import helpers_1
+
+
+class Model:
+    def predict(self, model_input: Any) -> Any:
+        return helpers_1.foo(123)

truss/tests/test_data/test_truss_with_error/packages/helpers_1.py

@@ -0,0 +1,5 @@
+import helpers_2
+
+
+def foo(x):
+    return helpers_2.bar(x)

truss/tests/test_data/test_truss_with_error/packages/helpers_2.py

@@ -0,0 +1,2 @@
+def bar(x):
+    raise Exception("Crashed in `bar`.")

truss/tests/test_docker.py

@@ -1,6 +1,7 @@
 import pytest
 from python_on_whales import docker
-from truss.docker import get_urls_from_container
+
+from truss.util.docker import get_urls_from_container
 
 
 @pytest.fixture