Make all non-optional

More comments and stuff

WIP

Fix more tests

Remove top-level symbols. Move more stuff

Pre-commit

fix constants/tests

Fix install commands and builder group

Manage poetry extras

Isolate base, cleanup trash

.github/workflows/_integration_test_shared.yml

@@ -13,7 +13,7 @@ jobs:
     steps:
       - uses: actions/checkout@v4
       - uses: ./.github/actions/setup-python/
-      - run: poetry install
+      - run: poetry install --with=dev,dev-server --extras=all
       - run: |
           poetry run pytest truss/tests \
           --durations=0 -m 'integration' \
@@ -34,7 +34,7 @@ jobs:
     steps:
       - uses: actions/checkout@v4
       - uses: ./.github/actions/setup-python/
-      - run: poetry install
+      - run: poetry install --with=dev,dev-server --extras=all
       - run: |
          poetry run pytest truss-chains/tests \
          --durations=0 -m 'integration' \

.github/workflows/integration-tests.yml

@@ -25,6 +25,7 @@ jobs:
           fetch-depth: 2
       - uses: ./.github/actions/detect-versions/
         id: versions
+
   build-and-push-truss-base-images-if-needed:
     needs: [detect-version-changed]
     if: needs.detect-version-changed.outputs.build_base_images == 'true'
@@ -46,7 +47,7 @@ jobs:
 
       - uses: actions/checkout@v4
       - uses: ./.github/actions/setup-python/
-      - run: poetry install
+      - run: poetry install --with=dev,dev-server --extras=all
       - shell: bash
         run: |
           poetry run bin/generate_base_images.py \

.github/workflows/main.yml

@@ -28,6 +28,7 @@ jobs:
           fetch-depth: 2
       - uses: ./.github/actions/detect-versions/
         id: versions
+
   build-and-push-truss-base-images-if-needed:
     needs: [detect-version-changed]
     if: needs.detect-version-changed.outputs.build_base_images == 'true'
@@ -49,7 +50,7 @@ jobs:
 
       - uses: actions/checkout@v4
       - uses: ./.github/actions/setup-python/
-      - run: poetry install
+      - run: poetry install --with=dev,dev-server --extras=all
       - shell: bash
         run: |
           poetry run bin/generate_base_images.py \
@@ -64,7 +65,7 @@ jobs:
     if: ${{ !failure() && !cancelled() && (needs.build-and-push-truss-base-images-if-needed.result == 'success' || needs.build-and-push-truss-base-images-if-needed.result == 'skipped') }}
     uses: ./.github/workflows/_integration_test_shared.yml
 
-  publish-to-pypi:
+  publish-rc-to-pypi:
     needs: [detect-version-changed]
     if: ${{ !failure() && !cancelled() && needs.detect-version-changed.outputs.release_version == 'true' && needs.detect-version-changed.outputs.is_prerelease_version == 'true' }}
     runs-on: ubuntu-20.04
@@ -84,7 +85,7 @@ jobs:
       - uses: ./.github/actions/setup-python/
 
       - name: Install poetry packages
-        run: poetry install --no-dev
+        run: poetry install --extras=all
 
       - name: Build
         run: poetry build

.github/workflows/pr.yml

@@ -15,7 +15,7 @@ jobs:
         with:
           lfs: true
       - uses: ./.github/actions/setup-python/
-      - run: poetry install
+      - run: poetry install --with=dev,dev-server --extras=all
       - run: poetry run pre-commit run --all-files
         env:
           SKIP: ruff
@@ -31,7 +31,7 @@ jobs:
         with:
           lfs: true
       - uses: ./.github/actions/setup-python/
-      - run: poetry install
+      - run: poetry install --with=dev,dev-server --extras=all
       - name: run tests
         run: poetry run pytest --durations=0  -m 'not integration' --junitxml=report.xml
       - name: Publish Test Report  # Not sure how to display this in the UI for non PRs.
@@ -50,7 +50,6 @@ jobs:
           use-verbose-mode: "yes"
           folder-path: "docs"
 
-
   enforce-chains-example-docs-sync:
     runs-on: ubuntu-20.04
     steps:

.github/workflows/release-truss-utils.yml

@@ -40,7 +40,7 @@ jobs:
 
       - name: Install poetry packages
         working-directory: truss-utils
-        run: poetry install --no-dev
+        run: poetry install
 
       - name: Build
         working-directory: truss-utils

.github/workflows/release.yml

@@ -60,7 +60,7 @@ jobs:
           tags: baseten/truss-context-builder:v${{ needs.detect-version-changed.outputs.new_version }}
           labels: ${{ steps.meta.outputs.labels }}
 
-  publish-to-pypi:
+  publish-release-to-pypi:
     needs: [detect-version-changed]
     if: ${{ !failure() && !cancelled() && needs.detect-version-changed.outputs.release_version == 'true' }}
     runs-on: ubuntu-20.04
@@ -80,7 +80,7 @@ jobs:
       - uses: ./.github/actions/setup-python/
 
       - name: Install poetry packages
-        run: poetry install --no-dev
+        run: poetry install --extras=all
 
       - name: Build
         run: poetry build

.pre-commit-config.yaml

@@ -22,6 +22,11 @@ repos:
       - id: ruff-format
   - repo: local
     hooks:
+      - id: pyproject.toml
+        name: pyproject.toml
+        entry: poetry run python bin/pyproject_toml_linter.py
+        language: python
+        additional_dependencies: [tomlkit]
       - id: mypy
         name: mypy-local
         entry: poetry run mypy

CONTRIBUTING.md

@@ -4,8 +4,8 @@ Truss was first created at [Baseten](https://baseten.co), but as an open and liv
 
 We use GitHub features for project management on Truss:
 
-* For bugs and feature requests, file an issue.
-* For changes and updates, create a pull request.
+* For bugs and feature requests, [file an issue](https://github.com/basetenlabs/truss/issues).
+* For changes and updates, create a [pull request](https://github.com/basetenlabs/truss/pulls).
 * To view and comment on the roadmap, [check the projects tab](https://github.com/orgs/basetenlabs/projects/3).
 
 ## Local development
@@ -30,7 +30,7 @@ asdf plugin add poetry
 asdf install
 
 # Install poetry dependencies
-poetry install
+poetry install --with=dev,dev-server --extras=all
 
 # And finally precommit
 poetry run pre-commit install

bin/codespace_post_create.sh

@@ -1,4 +1,4 @@
 #! /bin/bash
-poetry install
+poetry install --with=dev,dev-server --extras=all
 poetry run pre-commit install
 git lfs install

bin/generate_base_images.py

@@ -9,7 +9,7 @@
 from typing import List, Optional, Set
 
 from jinja2 import Environment, FileSystemLoader
-from truss.constants import SUPPORTED_PYTHON_VERSIONS
+from truss.base.constants import SUPPORTED_PYTHON_VERSIONS
 from truss.contexts.image_builder.util import (
     truss_base_image_name,
     truss_base_image_tag,

bin/pyproject_toml_linter.py

@@ -0,0 +1,77 @@
+# type: ignore  # tomlkit APIs are messy.
+import collections
+import pathlib
+from typing import DefaultDict, Set
+
+import tomlkit
+
+
+def _populate_extras(pyproject_path: pathlib.Path) -> None:
+    with pyproject_path.open("r", encoding="utf-8") as file:
+        original_content = file.read()
+        content = tomlkit.parse(original_content)
+
+    dependencies = content["tool"]["poetry"]["dependencies"]
+    dependency_metadata = content["tool"]["dependency_metadata"]
+
+    extra_sections: DefaultDict[str, Set[str]] = collections.defaultdict(set)
+    all_deps: Set[str] = set()
+
+    for key, value in dependencies.items():
+        if isinstance(value, dict):
+            is_optional = value.get("optional", False)
+        else:
+            is_optional = False  # Base dependencies.
+
+        if not is_optional:
+            continue
+
+        if key not in dependency_metadata:
+            raise ValueError(
+                f"`{key}` is missing in `[tool.dependency_metadata]`. "
+                f"(file: {pyproject_path}). Please add metadata."
+            )
+        metadata = dependency_metadata[key]
+        components = metadata["components"].split(",")
+        for component in components:
+            if component == "base":
+                continue
+            extra_sections[component].add(key)
+            all_deps.add(key)
+
+    for key in dependency_metadata.keys():
+        if key not in dependencies:
+            raise ValueError(
+                f"`{key}` in `[tool.dependency_metadata]` is not in "
+                "`[tool.poetry.dependencies]`. "
+                f"(file: {pyproject_path}). Please remove or sync."
+            )
+
+    extras_section = tomlkit.table()
+    for extra_section, deps in extra_sections.items():
+        extras_section[extra_section] = tomlkit.array()
+        extras_section[extra_section].extend(sorted(deps))
+
+    extras_section["all"] = tomlkit.array()
+    extras_section["all"].extend(sorted(all_deps))
+
+    if "extras" not in content["tool"]["poetry"]:
+        raise ValueError("Expected section [tool.poetry.extras] to be present.")
+
+    content["tool"]["poetry"]["extras"] = extras_section
+
+    updated_content = tomlkit.dumps(content)
+
+    # Compare the content before and after; if changes were made, fail the check
+    if original_content != updated_content:
+        with pyproject_path.open("w", encoding="utf-8") as file:
+            file.write(updated_content)
+        print(f"File '{pyproject_path}' was updated. Please re-stage the changes.")
+        exit(1)
+
+    print("No changes detected.")
+
+
+if __name__ == "__main__":
+    pyproject_file = pathlib.Path(__file__).parent.parent.resolve() / "pyproject.toml"
+    _populate_extras(pyproject_file)

context_builder.Dockerfile

@@ -24,5 +24,4 @@ COPY ./README.md ./README.md
 
 # https://python-poetry.org/docs/configuration/#virtualenvsin-project
 # to write to project root .venv file to be used for context builder test
-RUN poetry config virtualenvs.in-project true \
-    && poetry install --only builder
+RUN poetry config virtualenvs.in-project true && poetry install --extras=all

docs/chains/doc_gen/API-reference.mdx

@@ -281,12 +281,12 @@ For example, model weight caching can be used like this:
 
 ```python
 import truss_chains as chains
-from truss import truss_config
+from truss.base import truss_config
 
 mistral_cache = truss_config.ModelRepo(
     repo_id="mistralai/Mistral-7B-Instruct-v0.2",
     allow_patterns=["*.json", "*.safetensors", ".model"]
-  )
+)
 chains.Assets(cached=[mistral_cache], ...)
 ```