From 34602178d14820e80fae22ae41efb0f5fa29295e Mon Sep 17 00:00:00 2001 From: Marius Killinger <155577904+marius-baseten@users.noreply.github.com> Date: Tue, 19 Nov 2024 10:10:36 -0800 Subject: [PATCH] Fix model wrapper logger, and better XP around chains docker image validation error (#1251) --- docs/chains/doc_gen/API-reference.mdx | 16 +++---- docs/chains/doc_gen/reference.patch | 20 ++++---- truss-chains/truss_chains/definitions.py | 12 +++++ truss/templates/server/model_wrapper.py | 8 +++- truss/templates/server/truss_server.py | 5 +- truss/tests/test_model_inference.py | 60 +++++++++++++++++++++--- 6 files changed, 95 insertions(+), 26 deletions(-) diff --git a/docs/chains/doc_gen/API-reference.mdx b/docs/chains/doc_gen/API-reference.mdx index 6ff27670b..fe3d42a1f 100644 --- a/docs/chains/doc_gen/API-reference.mdx +++ b/docs/chains/doc_gen/API-reference.mdx @@ -214,14 +214,14 @@ modules and keep their requirement files right next their python source files. **Parameters:** -| Name | Type | Description | -|-------------------------|----------------------------------------------------------------------------------------------|--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| -| `base_image` | *[BasetenImage](#truss-chains-basetenimage)\|[CustomImage](#truss-chains-customimage)* | The base image used by the chainlet. Other dependencies and assets are included as additional layers on top of that image. You can choose a baseten default image for a supported python version (e.g. `BasetenImage.PY311`), this will also include GPU drivers if needed, or provide a custom image (e.g. `CustomImage(image="python:3.11-slim")`). | -| `pip_requirements_file` | *AbsPath\|None* | Path to a file containing pip requirements. The file content is naively concatenated with `pip_requirements`. | -| `pip_requirements` | *list[str]* | A list of pip requirements to install. The items are naively concatenated with the content of the `pip_requirements_file`. | -| `apt_requirements` | *list[str]* | A list of apt requirements to install. | -| `data_dir` | *AbsPath\|None* | Data from this directory is copied into the docker image and accessible to the remote chainlet at runtime. | -| `external_package_dirs` | *list[AbsPath]\|None* | A list of directories containing additional python packages outside the chain’s workspace dir, e.g. a shared library. This code is copied into the docker image and importable at runtime. | +| Name | Type | Description | +|-------------------------|----------------------------------------------------------------------------------------------------|--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| +| `base_image` | *[BasetenImage](#class-truss-chains-basetenimage)\|[CustomImage](#class-truss-chains-customimage)* | The base image used by the chainlet. Other dependencies and assets are included as additional layers on top of that image. You can choose a baseten default image for a supported python version (e.g. `BasetenImage.PY311`), this will also include GPU drivers if needed, or provide a custom image (e.g. `CustomImage(image="python:3.11-slim")`). | +| `pip_requirements_file` | *AbsPath\|None* | Path to a file containing pip requirements. The file content is naively concatenated with `pip_requirements`. | +| `pip_requirements` | *list[str]* | A list of pip requirements to install. The items are naively concatenated with the content of the `pip_requirements_file`. | +| `apt_requirements` | *list[str]* | A list of apt requirements to install. | +| `data_dir` | *AbsPath\|None* | Data from this directory is copied into the docker image and accessible to the remote chainlet at runtime. | +| `external_package_dirs` | *list[AbsPath]\|None* | A list of directories containing additional python packages outside the chain’s workspace dir, e.g. a shared library. This code is copied into the docker image and importable at runtime. | ### *class* `truss_chains.BasetenImage` diff --git a/docs/chains/doc_gen/reference.patch b/docs/chains/doc_gen/reference.patch index ecf741a4f..fd8d8fa29 100644 --- a/docs/chains/doc_gen/reference.patch +++ b/docs/chains/doc_gen/reference.patch @@ -1,5 +1,5 @@ ---- docs/chains/doc_gen/generated-reference.mdx 2024-11-07 16:51:32.687418306 -0800 -+++ docs/chains/doc_gen/API-reference.mdx 2024-11-07 16:58:24.661786055 -0800 +--- docs/chains/doc_gen/generated-reference.mdx 2024-11-14 15:10:37.862189314 -0800 ++++ docs/chains/doc_gen/API-reference.mdx 2024-11-18 12:04:23.725353699 -0800 @@ -24,31 +24,28 @@ dependency of another chainlet. The return value of `depends` is intended to be used as a default argument in a chainlet’s `__init__`-method. @@ -267,14 +267,14 @@ -#### pip_requirements *: list[str]* - -#### pip_requirements_file *: AbsPath | None* -+| Name | Type | Description | -+|-------------------------|----------------------------------------------------------------------------------------------|--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| -+| `base_image` | *[BasetenImage](#truss-chains-basetenimage)\|[CustomImage](#truss-chains-customimage)* | The base image used by the chainlet. Other dependencies and assets are included as additional layers on top of that image. You can choose a baseten default image for a supported python version (e.g. `BasetenImage.PY311`), this will also include GPU drivers if needed, or provide a custom image (e.g. `CustomImage(image="python:3.11-slim")`). | -+| `pip_requirements_file` | *AbsPath\|None* | Path to a file containing pip requirements. The file content is naively concatenated with `pip_requirements`. | -+| `pip_requirements` | *list[str]* | A list of pip requirements to install. The items are naively concatenated with the content of the `pip_requirements_file`. | -+| `apt_requirements` | *list[str]* | A list of apt requirements to install. | -+| `data_dir` | *AbsPath\|None* | Data from this directory is copied into the docker image and accessible to the remote chainlet at runtime. | -+| `external_package_dirs` | *list[AbsPath]\|None* | A list of directories containing additional python packages outside the chain’s workspace dir, e.g. a shared library. This code is copied into the docker image and importable at runtime. | ++| Name | Type | Description | ++|-------------------------|----------------------------------------------------------------------------------------------------|--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| ++| `base_image` | *[BasetenImage](#class-truss-chains-basetenimage)\|[CustomImage](#class-truss-chains-customimage)* | The base image used by the chainlet. Other dependencies and assets are included as additional layers on top of that image. You can choose a baseten default image for a supported python version (e.g. `BasetenImage.PY311`), this will also include GPU drivers if needed, or provide a custom image (e.g. `CustomImage(image="python:3.11-slim")`). | ++| `pip_requirements_file` | *AbsPath\|None* | Path to a file containing pip requirements. The file content is naively concatenated with `pip_requirements`. | ++| `pip_requirements` | *list[str]* | A list of pip requirements to install. The items are naively concatenated with the content of the `pip_requirements_file`. | ++| `apt_requirements` | *list[str]* | A list of apt requirements to install. | ++| `data_dir` | *AbsPath\|None* | Data from this directory is copied into the docker image and accessible to the remote chainlet at runtime. | ++| `external_package_dirs` | *list[AbsPath]\|None* | A list of directories containing additional python packages outside the chain’s workspace dir, e.g. a shared library. This code is copied into the docker image and importable at runtime. | ### *class* `truss_chains.BasetenImage` diff --git a/truss-chains/truss_chains/definitions.py b/truss-chains/truss_chains/definitions.py index 239147057..cdcb16272 100644 --- a/truss-chains/truss_chains/definitions.py +++ b/truss-chains/truss_chains/definitions.py @@ -191,6 +191,18 @@ class DockerImage(SafeModelNonSerializable): data_dir: Optional[AbsPath] = None external_package_dirs: Optional[list[AbsPath]] = None + @pydantic.root_validator(pre=True) + def migrate_fields(cls, values): + if "base_image" in values: + base_image = values["base_image"] + if isinstance(base_image, str): + doc_link = "https://docs.baseten.co/chains-reference/sdk#class-truss-chains-dockerimage" + raise ChainsUsageError( + "`DockerImage.base_image` as string is deprecated. Specify as " + f"`BasetenImage` or `CustomImage` (see docs: {doc_link})." + ) + return values + class ComputeSpec(pydantic.BaseModel): """Parsed and validated compute. See ``Compute`` for more information.""" diff --git a/truss/templates/server/model_wrapper.py b/truss/templates/server/model_wrapper.py index 0cba81bf5..82bab57d4 100644 --- a/truss/templates/server/model_wrapper.py +++ b/truss/templates/server/model_wrapper.py @@ -284,7 +284,13 @@ def __init__(self, config: Dict, tracer: sdk_trace.Tracer): self._tracer = tracer self._maybe_model = None self._maybe_model_descriptor = None - self._logger = logging.getLogger() + # We need a logger that has all our server JSON logging setup applied in its + # handlers and where this also hold in the loading thread. Creating a new + # instance does not carry over the setup into the thread and using unspecified + # `getLogger` may return non-compliant loggers if depdencies override the root + # logger (c.g. https://github.com/numpy/numpy/issues/24213). We chose to get + # the uvicorn logger that is setup in `truss_server`. + self._logger = logging.getLogger("uvicorn") self.name = MODEL_BASENAME self._load_lock = Lock() self._status = ModelWrapper.Status.NOT_READY diff --git a/truss/templates/server/truss_server.py b/truss/templates/server/truss_server.py index 657198dda..42b2293ae 100644 --- a/truss/templates/server/truss_server.py +++ b/truss/templates/server/truss_server.py @@ -270,7 +270,7 @@ async def _shutdown_if_load_fails(self): await asyncio.sleep(0.5) if self._model.load_failed: assert self._server is not None - logging.info("Trying shut down.") + logging.info("Trying shut down after failed model load.") self._server.should_exit = True return @@ -330,6 +330,9 @@ def start(self): if self._config["runtime"].get("enable_debug_logs", False) else "INFO" ) + # Warning: `ModelWrapper` depends on correctly setup `uvicorn` logger, + # if you change/remove that logger, make sure `ModelWrapper` has a suitable + # alternative logger that is also correctly setup in the load thread. cfg = uvicorn.Config( self.create_application(), # We hard-code the http parser as h11 (the default) in case the user has diff --git a/truss/tests/test_model_inference.py b/truss/tests/test_model_inference.py index 67ea3bf26..6bfcc005a 100644 --- a/truss/tests/test_model_inference.py +++ b/truss/tests/test_model_inference.py @@ -13,7 +13,7 @@ from concurrent.futures import ThreadPoolExecutor from pathlib import Path from threading import Thread -from typing import Iterator, Mapping +from typing import Iterator, Mapping, Optional import httpx import opentelemetry.trace.propagation.tracecontext as tracecontext @@ -40,22 +40,34 @@ def anyio_backend(): return "asyncio" -def _log_contains_error(line: dict, error: str, message: str): +def _log_contains_line( + line: dict, message: str, level: str, error: Optional[str] = None +): return ( - line["levelname"] == "ERROR" - and line["message"] == message - and error in line["exc_info"] + line["levelname"] == level + and message in line["message"] + and (error is None or error in line["exc_info"]) ) def _assert_logs_contain_error(logs: str, error: str, message=DEFAULT_LOG_ERROR): loglines = [json.loads(line) for line in logs.splitlines()] - assert any(_log_contains_error(line, error, message) for line in loglines), ( + assert any( + _log_contains_line(line, message, "ERROR", error) for line in loglines + ), ( f"Did not find expected error in logs.\nExpected error: {error}\n" f"Expected message: {message}\nActual logs:\n{loglines}" ) +def _assert_logs_contain(logs: str, message: str, level: str = "INFO"): + loglines = [json.loads(line) for line in logs.splitlines()] + assert any(_log_contains_line(line, message, level) for line in loglines), ( + f"Did not find expected logs.\n" + f"Expected message: {message}\nActual logs:\n{loglines}" + ) + + class _PropagatingThread(Thread): """ _PropagatingThread allows us to run threads and keep track of exceptions @@ -104,6 +116,42 @@ def test_map_to_supported_python_version(python_version, expected_python_version assert out_python_version == expected_python_version +@pytest.mark.integration +def test_model_load_logs(test_data_path): + model = """ + from typing import Optional + import logging + class Model: + def load(self): + logging.info(f"User Load Message") + + def predict(self, model_input): + return self.environment_name + """ + config = "model_name: init-environment-truss" + with ensure_kill_all(), _temp_truss(model, config) as tr: + container = tr.docker_run( + local_port=8090, detach=True, wait_for_server_ready=True + ) + logs = container.logs() + _assert_logs_contain( + logs, + message="Executing model.load()", + ) + _assert_logs_contain( + logs, + message="Loading truss model from file", + ) + _assert_logs_contain( + logs, + message="Completed model.load()", + ) + _assert_logs_contain( + logs, + message="User Load Message", + ) + + @pytest.mark.integration def test_model_load_failure_truss(test_data_path): with ensure_kill_all():