Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix model wrapper logger, and better XP around chains docker image va… #1251

Merged
merged 1 commit into from
Nov 19, 2024
Merged

Fix model wrapper logger, and better XP around chains docker image va… #1251

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 8 additions & 8 deletions docs/chains/doc_gen/API-reference.mdx
View file
Open in desktop
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
Open in desktop
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
Open in desktop
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
Open in desktop
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
Open in desktop
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
Loading