Skip to content

Commit

Permalink
Fix model wrapper logger, and better XP around chains docker image va…
Browse files Browse the repository at this point in the history 
…lidation error
  • Loading branch information
@marius-baseten
marius-baseten committed Nov 18, 2024
1 parent 251dd2c commit f83eb65
Show file tree
Hide file tree
Showing 6 changed files with 95 additions and 26 deletions.
16 changes: 8 additions & 8 deletions docs/chains/doc_gen/API-reference.mdx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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`

Expand Down
20 changes: 10 additions & 10 deletions docs/chains/doc_gen/reference.patch
View file
Original file line number Diff line number Diff line change
@@ -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.
Expand Down Expand Up @@ -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`

Expand Down
12 changes: 12 additions & 0 deletions truss-chains/truss_chains/definitions.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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."""
Expand Down
8 changes: 7 additions & 1 deletion truss/templates/server/model_wrapper.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
5 changes: 4 additions & 1 deletion truss/templates/server/truss_server.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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
Expand Down
Loading

0 comments on commit f83eb65

Please sign in to comment.