[TaT] Modular Package Distribution.

bin/pyproject_toml_linter.py

@@ -6,10 +6,9 @@
 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)
+def _populate_extras(pyproject_path: pathlib.Path) -> bool:
+    original_content = pyproject_path.read_text()
+    content = tomlkit.loads(original_content)
 
     dependencies = content["tool"]["poetry"]["dependencies"]
     dependency_metadata = content["tool"]["dependency_metadata"]
@@ -64,14 +63,40 @@ def _populate_extras(pyproject_path: pathlib.Path) -> None:
 
     # 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)
+        pyproject_path.write_text(updated_content)
         print(f"File '{pyproject_path}' was updated. Please re-stage the changes.")
-        exit(1)
+        return True
+
+    return False
+
+
+def _update_shim_version(
+    pyproject_path: pathlib.Path, shim_pyproject_path: pathlib.Path
+) -> bool:
+    truss_lib_pyproject = tomlkit.parse(pyproject_path.read_text())
+    version = truss_lib_pyproject["tool"]["poetry"]["version"]
 
-    print("No changes detected.")
+    truss_shim_pyproject = tomlkit.parse(shim_pyproject_path.read_text())
+    if truss_shim_pyproject["tool"]["poetry"]["version"] == version:
+        return False
+
+    truss_shim_pyproject["tool"]["poetry"]["version"] = tomlkit.string(version)
+    truss_shim_pyproject["tool"]["poetry"]["dependencies"]["truss_lib"]["version"] = (
+        version
+    )
+    shim_pyproject_path.write_text(tomlkit.dumps(truss_shim_pyproject))
+    return True
 
 
 if __name__ == "__main__":
     pyproject_file = pathlib.Path(__file__).parent.parent.resolve() / "pyproject.toml"
-    _populate_extras(pyproject_file)
+    changed = _populate_extras(pyproject_file)
+    shim_pyproject_file = (
+        pathlib.Path(__file__).parent.parent.resolve() / "truss-shim" / "pyproject.toml"
+    )
+    changed |= _update_shim_version(pyproject_file, shim_pyproject_file)
+
+    if changed:
+        exit(1)
+    else:
+        print("No changes detected.")

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
-name = "truss"
-version = "0.9.64"
+name = "truss-lib"
+version = "0.10.0rc000"
 description = "A seamless bridge from model development to model delivery"
 license = "MIT"
 readme = "README.md"
@@ -53,7 +53,7 @@ truss-docker-build-setup = "truss.contexts.docker_build_setup:docker_build_setup
 #   use the following installation command `poetry install --with=dev,dev-server --extras=all`.
 # * All other dependencies are in the main group `tool.poetry.dependencies`. Base dependencies are
 #   at the top and non-optional.
-# * Dependencies from other compoents are listed after, and marked with `optional = false`. If a
+# * Dependencies from other compoents are listed after, and marked with `optional = true`. If a
 #   dependency is needed by mutlipe extras, only add it once, but see next step. This also ensures
 #   that poetry resolves *all* dependencies from all extras to be globally consistent.
 # * Since poetry groups don't work with extras, we need to make the association between a dependency
@@ -70,66 +70,63 @@ truss-docker-build-setup = "truss.contexts.docker_build_setup:docker_build_setup
 #  isolated.
 [tool.poetry.dependencies]
 # "base" dependencies.
-# "When using chains, 3.9 will be required at runtime, but other truss functionality works with 3.8.
 python = ">=3.9,<3.13"
 huggingface_hub = ">=0.25.0"
 pydantic = ">=1.10.0"  # We cannot upgrade to v2, due to customer constraints.
 PyYAML = ">=6.0"
 single-source = "^0.3.0"
 # "non-base" dependencies.
-# TODO: until we have resolved the question on how users can install the local tools frictionless
-#  (extras cannot be marked to be included by default), all below packages are non-optional.
-#  This also means that so far extras defined in `[tool.poetry.extras]` don't have any meaning,
-#  since everything is globally included anyway.
-Jinja2 = { version = "^3.1.2", optional = false }
-aiofiles = { version = "^24.1.0", optional = false }
-aiohttp = { version = "^3.10.10", optional = false }
-blake3 = { version = "^0.3.3", optional = false }
-boto3 = { version = "^1.34.85", optional = false }
-click = { version = "^8.0.3", optional = false }
-fastapi = { version =">=0.109.1", optional = false }
-google-cloud-storage = { version = "2.10.0", optional = false }
-httpx = { version = ">=0.24.1", optional = false }
-inquirerpy = { version = "^0.3.4", optional = false }
-libcst = { version = "<1.2.0", optional = false }
-loguru = { version = ">=0.7.2", optional = false }
-packaging = { version = ">=20.9", optional = false }
-pathspec = { version = ">=0.9.0", optional = false }
-psutil = { version = ">=5.9.4", optional = false }
-python-json-logger = { version = ">=2.0.2", optional = false }
-python-on-whales = { version = "^0.68.0", optional = false }
-requests = { version = ">=2.31", optional = false }
-rich = { version = "^13.4.2", optional = false }
-rich-click = { version = "^1.6.1", optional = false }
-ruff = { version = ">=0.4.8", optional = false }  # Not a dev dep, needed for chains code gen.
-tenacity = { version = "^8.0.1", optional = false }
-watchfiles = { version = "^0.19.0", optional = false }
+Jinja2 = { version = "^3.1.2", optional = true }
+aiofiles = { version = "^24.1.0", optional = true }
+aiohttp = { version = "^3.10.10", optional = true }
+blake3 = { version = "^0.3.3", optional = true }
+boto3 = { version = "^1.34.85", optional = true }
+click = { version = "^8.0.3", optional = true }
+fastapi = { version =">=0.109.1", optional = true }
+google-cloud-storage = { version = "2.10.0", optional = true }
+httpx = { version = ">=0.24.1", optional = true }
+inquirerpy = { version = "^0.3.4", optional = true }
+libcst = { version = "<1.2.0", optional = true }
+loguru = { version = ">=0.7.2", optional = true }
+packaging = { version = ">=20.9", optional = true }
+pathspec = { version = ">=0.9.0", optional = true }
+psutil = { version = ">=5.9.4", optional = true }
+python-json-logger = { version = ">=2.0.2", optional = true }
+python-on-whales = { version = "^0.68.0", optional = true }
+requests = { version = ">=2.31", optional = true }
+rich = { version = "^13.4.2", optional = true }
+rich-click = { version = "^1.6.1", optional = true }
+ruff = { version = ">=0.4.8", optional = true }  # Not a dev dep, needed for chains code gen.
+tenacity = { version = "^8.0.1", optional = true }
+watchfiles = { version = "^0.19.0", optional = true }
 
 
 [tool.dependency_metadata]
 # `base` / `main` deps which are non-optional are always included and don't need to be added here.
-Jinja2 = { components = "other" }
-aiofiles = { components = "other" }
-aiohttp = { components = "other" }
-blake3 = { components = "other" }
-boto3 = { components = "other" }
-click = { components = "other" }
-google-cloud-storage = { components = "other" }
-httpx = { components = "other" }
-inquirerpy = { components = "other" }
-libcst = { components = "other" }
-loguru = { components = "other" }
-packaging = { components = "other" }
-pathspec = { components = "other" }
-psutil = { components = "other" }
-python-json-logger = { components = "other" }
-python-on-whales = { components = "other" }
-requests = { components = "other" }
-rich = { components = "other" }
-rich-click = { components = "other" }
-ruff = { components = "other" }
-tenacity = { components = "other" }
-watchfiles = { components = "other" }
+Jinja2 = { components = "other,cli" }
+aiofiles = { components = "other,cli" }
+aiohttp = { components = "other,cli" }
+blake3 = { components = "other,cli" }
+boto3 = { components = "other,cli" }
+click = { components = "other,cli" }
+fastapi = { components = "other,cli" }
+google-cloud-storage = { components = "other,cli" }
+httpx = { components = "other,cli" }
+inquirerpy = { components = "other,cli" }
+libcst = { components = "other,cli" }
+loguru = { components = "other,cli" }
+packaging = { components = "other,cli" }
+pathspec = { components = "other,cli" }
+psutil = { components = "other,cli" }
+python-json-logger = { components = "other,cli" }
+python-on-whales = { components = "other,cli" }
+requests = { components = "other,cli" }
+ruff = { components = "other,cli" }
+tenacity = { components = "other,cli" }
+watchfiles = { components = "other,cli" }
+# CLI only
+rich = { components = "cli" }
+rich-click = { components = "cli" }
 
 [tool.poetry.group.dev.dependencies]
 # These packages are needed as the dev/testing tooling
@@ -180,6 +177,7 @@ requires = ["poetry-core>=1.2.1"]
 ignore_missing_imports = true
 python_version = "3.9"
 plugins = ["pydantic.mypy"]
+explicit_package_bases = true
 
 [tool.pytest.ini_options]
 markers = [
@@ -231,4 +229,6 @@ max-line-length = 120
 
 # Note: `tool.poetry.extras` was autogenerated by `pyproject_toml_linter.py`, do not edit manually.
 [tool.poetry.extras]
-all = []
+other = ["Jinja2", "aiofiles", "aiohttp", "blake3", "boto3", "click", "fastapi", "google-cloud-storage", "httpx", "inquirerpy", "libcst", "loguru", "packaging", "pathspec", "psutil", "python-json-logger", "python-on-whales", "requests", "ruff", "tenacity", "watchfiles"]
+cli = ["Jinja2", "aiofiles", "aiohttp", "blake3", "boto3", "click", "fastapi", "google-cloud-storage", "httpx", "inquirerpy", "libcst", "loguru", "packaging", "pathspec", "psutil", "python-json-logger", "python-on-whales", "requests", "rich", "rich-click", "ruff", "tenacity", "watchfiles"]
+all = ["Jinja2", "aiofiles", "aiohttp", "blake3", "boto3", "click", "fastapi", "google-cloud-storage", "httpx", "inquirerpy", "libcst", "loguru", "packaging", "pathspec", "psutil", "python-json-logger", "python-on-whales", "requests", "rich", "rich-click", "ruff", "tenacity", "watchfiles"]

truss-shim/pyproject.toml

@@ -0,0 +1,23 @@
+[tool.poetry]
+name = "truss-shim"  # TODO: make atomic switch of `truss` -> `truss-lib` & `truss-shim` -> `truss`.
+version = "0.10.0rc000"
+description = "A seamless bridge from model development to model delivery"
+authors = ["Pankaj Gupta <[email protected]>", "Phil Howes <[email protected]>"]
+license = "MIT"
+
+packages = [
+    { include = "truss", from = "."}
+]
+
+[tool.poetry.scripts]
+truss = "truss.cli.cli:truss_cli"
+
+[tool.poetry.dependencies]
+# When using chains, 3.9 will be required at runtime, but other truss functionality works with 3.8.
+python = ">=3.8,<3.13"
+# TODO: carve out `cli` extra selectively instead of `all`.
+truss_lib = { version = "0.10.0rc000", extras = ["all"] }
+
+[build-system]
+build-backend = "poetry.core.masonry.api"
+requires = ["poetry-core>=1.2.1"]

truss-shim/truss/__init__.py

@@ -0,0 +1,8 @@
+from truss import version
+from truss.api.api import login, push, whoami
+from truss.base import truss_config
+from truss.truss_handle.build import load  # TODO: Refactor all usages and remove.
+
+__version__ = version.VERSION
+
+__all__ = ["push", "login", "load", "whoami", "truss_config"]

truss/cli/cli.py

@@ -16,10 +16,10 @@
 import rich.table
 import rich.traceback
 import rich_click as click
-import truss
 from InquirerPy import inquirer
 from rich import progress
 from rich.console import Console
+from truss import version
 from truss.base.constants import (
     PRODUCTION_ENVIRONMENT_NAME,
     TRTLLM_MIN_MEMORY_REQUEST_GI,
@@ -86,6 +86,9 @@
     ]
 }
 
+VERSION = version.VERSION
+
+
 console = Console()
 
 error_console = Console(stderr=True, style="bold red")
@@ -179,7 +182,7 @@ def print_help() -> None:
 
 @click.group(name="truss", invoke_without_command=True)  # type: ignore
 @click.pass_context
-@click.version_option(truss.version())
+@click.version_option(VERSION)
 @log_level_option
 def truss_cli(ctx) -> None:
     """truss: The simplest way to serve models in production"""
@@ -341,7 +344,7 @@ def run(
 )
 @error_handling
 def login(api_key: Optional[str]):
-    from truss.api import login
+    from truss.api.api import login
 
     if not api_key:
         remote_config = inquire_remote_config()
@@ -362,7 +365,7 @@ def whoami(remote: Optional[str]):
     """
     Shows user information and exit.
     """
-    from truss.api import whoami
+    from truss.api.api import whoami
 
     if not remote:
         remote = inquire_remote_name(RemoteFactory.get_available_config_names())

truss/contexts/image_builder/util.py

@@ -1,7 +1,7 @@
 from pathlib import Path
 from typing import Optional
 
-from truss import __version__
+from truss.version import VERSION
 
 # This needs to be updated whenever we want to update
 # base images on a merge. Updating this version will cause
@@ -39,7 +39,7 @@ def truss_base_image_tag(
     python_version: str, use_gpu: bool, version_tag: Optional[str] = None
 ) -> str:
     if version_tag is None:
-        version_tag = f"v{__version__}"
+        version_tag = f"v{VERSION}"
 
     base_tag = python_version
     if use_gpu:

truss/remote/baseten/core.py

@@ -211,9 +211,9 @@ def get_dev_version_from_versions(versions: List[dict]) -> Optional[dict]:
         The version in versions corresponding to the development model, or None
         if no development model exists
     """
-    for version in versions:
-        if version["is_draft"] is True:
-            return version
+    for model_version in versions:
+        if model_version["is_draft"] is True:
+            return model_version
     return None
 
 
@@ -267,9 +267,9 @@ def get_prod_version_from_versions(versions: List[dict]) -> Optional[dict]:
     # Loop over versions instead of using the primary_version field because
     # primary_version is set to the development version ID if no published
     # models exist.
-    for version in versions:
-        if version["is_primary"]:
-            return version
+    for model_version in versions:
+        if model_version["is_primary"]:
+            return model_version
     return None
 
 

truss/remote/remote_factory.py

@@ -14,7 +14,7 @@
 from pathlib import Path
 from typing import Dict, List, Type
 
-from truss.remote.baseten import BasetenRemote
+from truss.remote.baseten.remote import BasetenRemote
 from truss.remote.truss_remote import RemoteConfig, TrussRemote
 
 USER_TRUSSRC_PATH = Path(os.environ.get("USER_TRUSSRC_PATH", "~/.trussrc")).expanduser()

truss/version.py

@@ -0,0 +1,5 @@
+from pathlib import Path
+
+from single_source import get_version
+
+VERSION = get_version(__name__, Path(__file__).parent.parent)