truss 0.9.12__tar.gz → 0.9.12.dev2__tar.gz

{truss-0.9.12 → truss-0.9.12.dev2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: truss
-Version: 0.9.12
+Version: 0.9.12.dev2
 Summary: A seamless bridge from model development to model delivery
 Home-page: https://github.com/basetenlabs/truss
 License: MIT

{truss-0.9.12 → truss-0.9.12.dev2}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "truss"
-version = "0.9.12"
+version = "0.9.12dev2"
 description = "A seamless bridge from model development to model delivery"
 license = "MIT"
 readme = "README.md"
@@ -117,7 +117,7 @@ src_paths = ["isort", "test"]
 [tool.mypy]
 ignore_missing_imports = true
 python_version = "3.9"
-plugins = ["pydantic.mypy"]
+
 
 [tool.pytest.ini_options]
 markers = [

{truss-0.9.12 → truss-0.9.12.dev2}/truss/cli/cli.py

@@ -635,7 +635,7 @@ def predict(
         "specifying, the command will not complete until the deployment is complete."
     ),
 )
-@error_handling
+# @error_handling
 def push(
     target_directory: str,
     remote: str,
@@ -670,16 +670,6 @@ def push(
         tr.spec.config.model_name = model_name
         tr.spec.config.write_to_yaml_file(tr.spec.config_path, verbose=False)
 
-    # Log a warning if using secrets without --trusted.
-    # TODO(helen): this could be moved to a separate function that includes more
-    #  config checks.
-    if tr.spec.config.secrets and not trusted:
-        not_trusted_text = (
-            "Warning: your Truss has secrets but was not pushed with --trusted. "
-            "Please push with --trusted to grant access to secrets."
-        )
-        console.print(not_trusted_text, style="red")
-
     # TODO(Abu): This needs to be refactored to be more generic
     service = remote_provider.push(
         tr,
@@ -708,6 +698,16 @@ def push(
 
         click.echo(draft_model_text)
 
+    # Log a warning if using secrets without --trusted.
+    # TODO(helen): this could be moved to a separate function that includes more
+    #  config checks.
+    if tr.spec.config.secrets and not trusted:
+        not_trusted_text = (
+            "Warning: your Truss has secrets but was not pushed with --trusted."
+            "Please push with --trusted to grant access to secrets."
+        )
+        console.print(not_trusted_text, style="red")
+
     if promote:
         promotion_text = (
             "Your Truss has been deployed as a production model. After it successfully "

{truss-0.9.12 → truss-0.9.12.dev2}/truss/config/trt_llm.py

@@ -29,7 +29,7 @@ class TrussTRTLLMQuantizationType(str, Enum):
 class TrussTRTLLMPluginConfiguration(BaseModel):
     multi_block_mode: bool = False
     paged_kv_cache: bool = True
-    gemm_plugin: str = "float16"
+    use_fused_mlp: bool = False
 
 
 class CheckpointSource(Enum):

{truss-0.9.12 → truss-0.9.12.dev2}/truss/constants.py

@@ -104,6 +104,6 @@ BASE_TRTLLM_REQUIREMENTS = [
     "tritonclient[all]==2.42.0",
     "transformers==4.40.2",
     "jinja2==3.1.3",
-    "truss==0.9.12",
+    "truss==0.9.11",
 ]
 OPENAI_COMPATIBLE_TAG = "openai-compatible"

{truss-0.9.12 → truss-0.9.12.dev2}/truss/remote/baseten/service.py

@@ -89,7 +89,7 @@ class BasetenService(TrussService):
             # will be a json with an `error` key.
             return parsed_response
 
-        return response.json()
+        return response.json()["model_output"]
 
     def authenticate(self) -> dict:
         return self._auth_service.authenticate().header()

{truss-0.9.12 → truss-0.9.12.dev2}/truss/test_data/test_basic_truss/config.yaml

@@ -9,7 +9,5 @@ resources:
   cpu: '1'
   memory: 2Gi
   use_gpu: false
-secrets:
-  foo: "foo"
-  bar: "bar"
+secrets: {}
 system_packages: []

{truss-0.9.12 → truss-0.9.12.dev2}/truss/tests/test_config.py

@@ -315,57 +315,29 @@ def test_empty_config():
 
 
 def test_from_yaml():
+    yaml_path = Path("test.yaml")
     data = {"description": "this is a test"}
-    with tempfile.NamedTemporaryFile(mode="w", delete=False) as yaml_file:
-        yaml_path = Path(yaml_file.name)
+    with yaml_path.open("w") as yaml_file:
         yaml.safe_dump(data, yaml_file)
 
-        result = TrussConfig.from_yaml(yaml_path)
+    result = TrussConfig.from_yaml(yaml_path)
+
+    assert result.description == "this is a test"
 
-        assert result.description == "this is a test"
+    yaml_path.unlink()
 
 
 def test_from_yaml_empty():
+    yaml_path = Path("test.yaml")
     data = {}
-    with tempfile.NamedTemporaryFile(mode="w", delete=False) as yaml_file:
-        yaml_path = Path(yaml_file.name)
+    with yaml_path.open("w") as yaml_file:
         yaml.safe_dump(data, yaml_file)
 
-        result = TrussConfig.from_yaml(yaml_path)
-
-        # test some attributes (should be default)
-        assert result.description is None
-        assert result.spec_version == "2.0"
-        assert result.bundled_packages_dir == "packages"
-
-
-def test_from_yaml_secrets_as_list():
-    data = {"description": "this is a test", "secrets": ["foo", "bar"]}
-    with tempfile.NamedTemporaryFile(mode="w", delete=False) as yaml_file:
-        yaml_path = Path(yaml_file.name)
-        yaml.safe_dump(data, yaml_file)
-
-        with pytest.raises(ValueError):
-            TrussConfig.from_yaml(yaml_path)
-
-
-def test_from_yaml_python_version():
-    yaml_path = Path("test.yaml")
-    invalid_py_version_data = {
-        "description": "this is a test",
-        "python_version": "py38",
-    }
-    with tempfile.NamedTemporaryFile(mode="w", delete=False) as yaml_file:
-        yaml_path = Path(yaml_file.name)
-        yaml.safe_dump(invalid_py_version_data, yaml_file)
-
-        with pytest.raises(ValueError):
-            TrussConfig.from_yaml(yaml_path)
+    result = TrussConfig.from_yaml(yaml_path)
 
-    valid_py_version_data = {"description": "this is a test", "python_version": "py39"}
-    with tempfile.NamedTemporaryFile(mode="w", delete=False) as yaml_file:
-        yaml_path = Path(yaml_file.name)
-        yaml.safe_dump(valid_py_version_data, yaml_file)
+    # test some attributes (should be default)
+    assert result.description is None
+    assert result.spec_version == "2.0"
+    assert result.bundled_packages_dir == "packages"
 
-        result = TrussConfig.from_yaml(yaml_path)
-        result.python_version == "py39"
+    yaml_path.unlink()

{truss-0.9.12 → truss-0.9.12.dev2}/truss/truss_config.py

@@ -39,8 +39,6 @@ DEFAULT_USE_GPU = False
 
 DEFAULT_BLOB_BACKEND = HTTP_PUBLIC_BLOB_BACKEND
 
-VALID_PYTHON_VERSIONS = ["py39", "py310", "py311"]
-
 # Set up logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -578,21 +576,6 @@ class TrussConfig:
         return TrussConfig.from_dict(self.to_dict())
 
     def validate(self):
-        if self.python_version not in VALID_PYTHON_VERSIONS:
-
-            raise ValueError(
-                f"Please ensure that `python_version` is one of {VALID_PYTHON_VERSIONS}"
-            )
-
-        if not isinstance(self.secrets, dict):
-            raise ValueError(
-                "Please ensure that `secrets` is a mapping of the form:\n"
-                "```\n"
-                "secrets:\n"
-                '  secret1: "some default value"\n'
-                '  secret2: "some other default value"\n'
-                "```"
-            )
         for secret_name in self.secrets:
             validate_secret_name(secret_name)
 

{truss-0.9.12 → truss-0.9.12.dev2}/truss/validation.py

@@ -1,5 +1,5 @@
 import re
-from pathlib import PosixPath
+from pathlib import Path
 from typing import Pattern, Set
 
 from truss.constants import REGISTRY_BUILD_SECRET_PREFIX
@@ -81,10 +81,9 @@ def _is_numeric(number_like: str) -> bool:
 def validate_python_executable_path(path: str) -> None:
     """
     This python executable path determines the python executable
-    used to run the inference server - check to see that it is an absolute path.
-    We use PosixPath -- TrussServer always runs on a Posix machine.
+    used to run the inference server - check to see that it is an absolute path
     """
-    if path and not PosixPath(path).is_absolute():
+    if path and not Path(path).is_absolute():
         raise ValidationError(
             f"Invalid relative python executable path {path}. Provide an absolute path"
         )

{truss-0.9.12 → truss-0.9.12.dev2}/truss-chains/truss_chains/__init__.py

@@ -23,10 +23,10 @@ from truss_chains.definitions import (
 )
 from truss_chains.public_api import (
     ChainletBase,
-    depends,
-    depends_context,
     deploy_remotely,
     entrypoint,
+    provide,
+    provide_context,
     run_local,
 )
 from truss_chains.stub import StubBase

{truss-0.9.12 → truss-0.9.12.dev2}/truss-chains/truss_chains/code_gen.py

@@ -131,7 +131,7 @@ def _gen_chainlet_import_and_ref(
 
 def _endpoint_signature_src(endpoint: definitions.EndpointAPIDescriptor) -> _Source:
     """
-    E.g.: `async def run_remote(self, data: str, num_partitions: int) -> tuple[list, int]:`
+    E.g.: `async def run(self, data: str, num_partitions: int) -> tuple[list, int]:`
     """
     if endpoint.is_generator:
         # TODO: implement generator.
@@ -229,7 +229,7 @@ def _gen_stub_src(chainlet: definitions.ChainletAPIDescriptor) -> _Source:
         def __init__(self, url: str, api_key: str) -> None:
             self._remote = stub.BasetenSession(url, api_key)
 
-        async def run_remote(self, data: str, num_partitions: int) -> tuple[SplitTextOutput, int]:
+        async def run(self, data: str, num_partitions: int) -> tuple[SplitTextOutput, int]:
             json_args = {"inputs": inputs.dict(), "extra_arg": extra_arg}
             json_result = await self._remote.predict_async(json_args)
             return (SplitTextOutput.parse_obj(json_result[0]), json_result[1])
@@ -317,13 +317,10 @@ def _gen_load_src(chainlet_descriptor: definitions.ChainletAPIDescriptor) -> _So
         # `dep.name` is the class name, while `name` is the argument name.
         stub_args.append(f"{name}=stub.factory({dep.name}, self._context)")
 
-    if chainlet_descriptor.has_context:
-        if stub_args:
-            init_args = f"{', '.join(stub_args)}, context=self._context"
-        else:
-            init_args = "context=self._context"
+    if stub_args:
+        init_args = f"context=self._context, {', '.join(stub_args)}"
     else:
-        init_args = ", ".join(stub_args)
+        init_args = "context=self._context"
 
     user_chainlet_ref = _gen_chainlet_import_and_ref(chainlet_descriptor)
     imports.update(user_chainlet_ref.imports)
@@ -389,6 +386,7 @@ def _gen_predict_src(chainlet_descriptor: definitions.ChainletAPIDescriptor) ->
 def _gen_truss_chainlet_model(
     chainlet_descriptor: definitions.ChainletAPIDescriptor,
 ) -> _Source:
+    logging.info(f"Generating truss model for chainlet `{chainlet_descriptor.name}`.")
     skeleton_tree = libcst.parse_module(
         pathlib.Path(model_skeleton.__file__).read_text()
     )
@@ -521,11 +519,10 @@ def _make_truss_config(
     # Compute.
     compute = chains_config.get_compute_spec()
     config.resources.cpu = str(compute.cpu_count)
-    config.resources.memory = str(compute.memory)
     config.resources.accelerator = compute.accelerator
     config.resources.use_gpu = bool(compute.accelerator.count)
     # TODO: expose this setting directly.
-    config.runtime.predict_concurrency = compute.predict_concurrency
+    config.runtime.predict_concurrency = compute.cpu_count
     # Image.
     image = chains_config.docker_image
     config.base_image = truss_config.BaseImage(image=image.base_image)
@@ -570,7 +567,6 @@ def gen_truss_chainlet(
     gen_root: pathlib.Path,
 ) -> pathlib.Path:
     # TODO: support file-based config (and/or merge file and python-src config values).
-    logging.info(f"Generating truss chainlet model for `{chainlet_descriptor.name}`.")
     remote_config = chainlet_descriptor.chainlet_cls.remote_config
     chainlet_name = remote_config.name or chainlet_descriptor.name
     # Filter needed services and customize options.

{truss-0.9.12 → truss-0.9.12.dev2}/truss-chains/truss_chains/definitions.py

@@ -10,15 +10,12 @@ from typing import (
     ClassVar,
     Generic,
     Iterable,
-    Literal,
     Mapping,
     Optional,
-    Protocol,
     Type,
     TypeVar,
     Union,
     cast,
-    runtime_checkable,
 )
 
 import pydantic
@@ -27,37 +24,22 @@ from truss import truss_config
 from truss.remote import baseten as baseten_remote
 from truss.remote import remote_cli, remote_factory
 
-UserConfigT = TypeVar("UserConfigT", bound=Union[pydantic.BaseModel, None])
+UserConfigT = TypeVar("UserConfigT", bound=Optional[pydantic.BaseModel])
 
 BASETEN_API_SECRET_NAME = "baseten_chain_api_key"
 SECRET_DUMMY = "***"
 TRUSS_CONFIG_CHAINS_KEY = "chains_metadata"
 
 GENERATED_CODE_DIR = ".chains_generated"
+PREDICT_ENDPOINT_NAME = "/predict"
 
 # Below arg names must correspond to `definitions.ABCChainlet`.
-ENDPOINT_METHOD_NAME = "run_remote"  # Chainlet method name exposed as endpoint.
+ENDPOINT_METHOD_NAME = "run"  # Referring to Chainlet method name exposed as endpoint.
 CONTEXT_ARG_NAME = "context"  # Referring to Chainlets `__init__` signature.
 SELF_ARG_NAME = "self"
 REMOTE_CONFIG_NAME = "remote_config"
 
 
-K = TypeVar("K", contravariant=True)
-V = TypeVar("V", covariant=True)
-
-
-@runtime_checkable
-class MappingNoIter(Protocol[K, V]):
-    def __getitem__(self, key: K) -> V:
-        ...
-
-    def __len__(self) -> int:
-        ...
-
-    def __contains__(self, key: K) -> bool:
-        ...
-
-
 class SafeModel(pydantic.BaseModel):
     """Pydantic base model with reasonable config."""
 
@@ -134,15 +116,10 @@ class DockerImage(SafeModelNonSerializable):
 class ComputeSpec(pydantic.BaseModel):
     # TODO: this is not stable yet and might change or refer back to truss.
     cpu_count: int = 1
-    predict_concurrency: int = 1
     memory: str = "2Gi"
     accelerator: truss_config.AcceleratorSpec = truss_config.AcceleratorSpec()
 
 
-CpuCountT = Literal["cpu_count"]
-CPU_COUNT: CpuCountT = "cpu_count"
-
-
 class Compute:
     """Builder to create ComputeSpec."""
 
@@ -156,7 +133,6 @@ class Compute:
         memory: str = "2Gi",
         gpu: Union[str, truss_config.Accelerator, None] = None,
         gpu_count: int = 1,
-        predict_concurrency: Union[int, CpuCountT] = 1,
     ) -> None:
         accelerator = truss_config.AcceleratorSpec()
         if gpu:
@@ -165,17 +141,9 @@ class Compute:
             accelerator = truss_config.AcceleratorSpec(
                 accelerator=truss_config.Accelerator(gpu), count=gpu_count
             )
-        if predict_concurrency == CPU_COUNT:
-            predict_concurrency_int = cpu_count
-        else:
-            assert isinstance(predict_concurrency, int)
-            predict_concurrency_int = predict_concurrency
 
         self._spec = ComputeSpec(
-            cpu_count=cpu_count,
-            memory=memory,
-            accelerator=accelerator,
-            predict_concurrency=predict_concurrency_int,
+            cpu_count=cpu_count, memory=memory, accelerator=accelerator
         )
 
     def get_spec(self) -> ComputeSpec:
@@ -210,7 +178,7 @@ class Assets:
 class RemoteConfig(SafeModelNonSerializable):
     """Bundles config values needed to deploy a Chainlet."""
 
-    docker_image: DockerImage = DockerImage()
+    docker_image: DockerImage
     compute: Compute = Compute()
     assets: Assets = Assets()
     name: Optional[str] = None
@@ -243,9 +211,13 @@ class DeploymentContext(generics.GenericModel, Generic[UserConfigT]):
         extra = "forbid"
 
     data_dir: Optional[pathlib.Path] = None
-    user_config: UserConfigT
+    user_config: UserConfigT = pydantic.Field(default=None)
     chainlet_to_service: Mapping[str, ServiceDescriptor] = {}
-    secrets: MappingNoIter[str, str]
+    # secrets: Optional[secrets_resolver.Secrets] = None
+    # TODO: above type results in `truss.server.shared.secrets_resolver.Secrets`
+    #   due to the templating, at runtime the object passed will be from
+    #   `shared.secrets_resolver` and give pydantic validation error.
+    secrets: Any = None
 
     def get_service_descriptor(self, chainlet_name: str) -> ServiceDescriptor:
         if chainlet_name not in self.chainlet_to_service:
@@ -277,24 +249,30 @@ class DeploymentContext(generics.GenericModel, Generic[UserConfigT]):
 class TrussMetadata(generics.GenericModel, Generic[UserConfigT]):
     """Plugin for the truss config (in config["model_metadata"]["chains_metadata"])."""
 
-    user_config: UserConfigT
+    user_config: UserConfigT = pydantic.Field(default=None)
     chainlet_to_service: Mapping[str, ServiceDescriptor] = {}
 
 
-class ABCChainlet(abc.ABC):
-    remote_config: ClassVar[RemoteConfig] = RemoteConfig(docker_image=DockerImage())
+class ABCChainlet(Generic[UserConfigT], abc.ABC):
+    remote_config: ClassVar[RemoteConfig]
     default_user_config: ClassVar[Optional[pydantic.BaseModel]] = None
     _init_is_patched: ClassVar[bool] = False
 
-    @classmethod
-    def has_custom_init(cls) -> bool:
-        return cls.__init__ is not object.__init__
+    _context: DeploymentContext[UserConfigT]
+
+    @abc.abstractmethod
+    def __init__(self, context: DeploymentContext[UserConfigT]) -> None:
+        ...
 
     # Cannot add this abstract method to API, because we want to allow arbitrary
     # arg/kwarg names and specifying any function signature here would give type errors
     # @abc.abstractmethod
-    # def run_remote(self, *args, **kwargs) -> Any:
-    #     ...
+    # def run(self, *args, **kwargs) -> Any: ...
+
+    @property
+    @abc.abstractmethod
+    def user_config(self) -> UserConfigT:
+        ...
 
 
 class TypeDescriptor(SafeModelNonSerializable):
@@ -331,7 +309,6 @@ class DependencyDescriptor(SafeModelNonSerializable):
 class ChainletAPIDescriptor(SafeModelNonSerializable):
     chainlet_cls: Type[ABCChainlet]
     src_path: str
-    has_context: bool
     dependencies: Mapping[str, DependencyDescriptor]
     endpoint: EndpointAPIDescriptor
     user_config_type: TypeDescriptor

{truss-0.9.12 → truss-0.9.12.dev2}/truss-chains/truss_chains/deploy.py

@@ -149,7 +149,7 @@ class ChainService:
     def run_url(self) -> str:
         return self.get_entrypoint.predict_url
 
-    def run_remote(self, json: Dict) -> Any:
+    def run(self, json: Dict) -> Any:
         return self.get_entrypoint.predict(json)
 
     def get_info(self) -> list[tuple[str, str, str]]:

truss-0.9.12.dev2/truss-chains/truss_chains/example_chainlet.py

@@ -0,0 +1,37 @@
+import truss_chains as chains
+
+
+class DummyGenerateData(chains.ChainletBase):
+    remote_config = chains.RemoteConfig(docker_image=chains.DockerImage())
+
+    def run(self) -> str:
+        return "abc"
+
+
+# Nesting the classes is a hack to make it *appear* like SplitText is from a different
+# module.
+class shared_chainlet:
+    class DummySplitText(chains.ChainletBase):
+        remote_config = chains.RemoteConfig(docker_image=chains.DockerImage())
+
+        def run(self, data: str) -> list[str]:
+            return [data[:2], data[2:]]
+
+
+class DummyExample(chains.ChainletBase):
+    remote_config = chains.RemoteConfig(docker_image=chains.DockerImage())
+
+    def __init__(
+        self,
+        context: chains.DeploymentContext = chains.provide_context(),
+        data_generator: DummyGenerateData = chains.provide(DummyGenerateData),
+        splitter: shared_chainlet.DummySplitText = chains.provide(
+            shared_chainlet.DummySplitText
+        ),
+    ) -> None:
+        super().__init__(context)
+        self._data_generator = data_generator
+        self._data_splitter = splitter
+
+    def run(self) -> list[str]:
+        return self._data_splitter.run(self._data_generator.run())