Merge pull request #1631 from basetenlabs/bump-version-0.9.88

Release 0.9.88

Author: marius-baseten

docs/examples/frameworks/pytorch.mdx

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "truss"
-version = "0.9.87"
+version = "0.9.88"
 description = "A seamless bridge from model development to model delivery"
 license = "MIT"
 readme = "README.md"
@@ -102,6 +102,7 @@ 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 }
+tomlkit = { version = ">=0.13.2", optional = false }
 watchfiles = { version = "^0.19.0", optional = false }
 
 
@@ -127,6 +128,7 @@ rich = { components = "other" }
 rich-click = { components = "other" }
 ruff = { components = "other" }
 tenacity = { components = "other" }
+tomlkit = { components = "other" }
 watchfiles = { components = "other" }
 
 [tool.poetry.group.dev.dependencies]
@@ -146,7 +148,6 @@ pytest-check = "^2.4.1"
 pytest-cov = "^3.0.0"
 pytest-split = ">=0.10.0"
 requests-mock = ">=1.11.0"
-tomlkit = ">=0.12"
 types-PyYAML = "^6.0.12.12"
 types-aiofiles = ">=24.1.0"
 types-requests = "==2.31.0.2"

docs/examples/frameworks/sklearn.mdx

@@ -32,7 +32,7 @@
 from truss.remote.baseten import remote as b10_remote
 from truss.remote.baseten import service as b10_service
 from truss.truss_handle import truss_handle
-from truss.util import log_utils
+from truss.util import log_utils, user_config
 from truss.util import path as truss_path
 from truss_chains import framework, private_types, public_types
 from truss_chains.deployment import code_gen
@@ -472,7 +472,7 @@ def _create_baseten_chain(
         remote_factory.RemoteFactory.create(remote=baseten_options.remote),
     )
 
-    if remote_provider.include_git_info or baseten_options.include_git_info:
+    if user_config.settings.include_git_info or baseten_options.include_git_info:
         truss_user_env = b10_types.TrussUserEnv.collect_with_git_info(
             baseten_options.working_dir
         )

docs/examples/frameworks/tensorflow.mdx

@@ -0,0 +1,26 @@
+from truss.base import truss_config
+from truss_train import definitions
+
+deploy_checkpoint = definitions.DeployCheckpointsConfig(
+    compute=definitions.Compute(
+        accelerator=truss_config.AcceleratorSpec(
+            accelerator=truss_config.Accelerator.A10G, count=4
+        )
+    ),
+    runtime=definitions.DeployCheckpointsRuntime(
+        environment_variables={
+            "HF_TOKEN": definitions.SecretReference(name="hf_access_token")
+        }
+    ),
+    checkpoint_details=definitions.CheckpointDetails(
+        base_model_id="unsloth/gemma-3-1b-it",
+        checkpoints=[
+            definitions.Checkpoint(
+                id="checkpoint-24", name="checkpoint-24", training_job_id="lqz4pw4"
+            ),
+            definitions.Checkpoint(
+                id="checkpoint-42", name="checkpoint-42", training_job_id="lqz4pw4"
+            ),
+        ],
+    ),
+)

docs/examples/frameworks/xgboost.mdx

@@ -0,0 +1,41 @@
+from truss.base import truss_config
+from truss_train import definitions
+
+deploy_checkpoint = definitions.DeployCheckpointsConfig(
+    compute=definitions.Compute(
+        accelerator=truss_config.AcceleratorSpec(
+            accelerator=truss_config.Accelerator.A10G, count=4
+        )
+    ),
+    runtime=definitions.DeployCheckpointsRuntime(
+        environment_variables={
+            "HF_TOKEN": definitions.SecretReference(name="hf_access_token")
+        }
+    ),
+    checkpoint_details=definitions.CheckpointDetails(
+        base_model_id="unsloth/gemma-3-1b-it",
+        checkpoints=[
+            definitions.Checkpoint(
+                id="checkpoint-24", name="checkpoint-24", training_job_id="lqz4pw4"
+            ),
+            definitions.Checkpoint(
+                id="checkpoint-42", name="checkpoint-42", training_job_id="lqz4pw4"
+            ),
+        ],
+    ),
+)
+runtime_config = definitions.Runtime(
+    start_commands=["/bin/bash ./my-entrypoint.sh"],
+    environment_variables={
+        "FOO_VAR": "FOO_VAL",
+        "BAR_VAR": definitions.SecretReference(name="BAR_SECRET"),
+    },
+)
+
+training_job = definitions.TrainingJob(
+    image=definitions.Image(base_image="base-image"),
+    compute=definitions.Compute(node_count=1, cpu_count=4),
+    runtime_config=runtime_config,
+)
+
+first_project = definitions.TrainingProject(name="first-project", job=training_job)

poetry.lock

@@ -11,21 +11,21 @@ def test_import_requires_training_project():
     job_src = TEST_ROOT / "import" / "config_without_training_project.py"
     match = r"No `.+` was found."
     with pytest.raises(ValueError, match=match):
-        with loader.import_target(job_src):
+        with loader.import_training_project(job_src):
             pass
 
 
 def test_import_requires_single_training_project():
     job_src = TEST_ROOT / "import" / "config_with_multiple_training_projects.py"
     match = r"Multiple `.+`s were found."
     with pytest.raises(ValueError, match=match):
-        with loader.import_target(job_src):
+        with loader.import_training_project(job_src):
             pass
 
 
 def test_import_with_single_training_project():
     job_src = TEST_ROOT / "import" / "config_with_single_training_project.py"
-    with loader.import_target(job_src) as training_project:
+    with loader.import_training_project(job_src) as training_project:
         assert training_project.name == "first-project"
         assert training_project.job.compute.cpu_count == 4
 
@@ -34,5 +34,44 @@ def test_import_directory_fails():
     job_src = TEST_ROOT / "import"
     match = r"You must point to a python file"
     with pytest.raises(ImportError, match=match):
-        with loader.import_target(job_src):
+        with loader.import_training_project(job_src):
             pass
+
+
+def test_import_deploy_checkpoints_config():
+    job_src = TEST_ROOT / "import" / "deploy_checkpoints_config.py"
+    with loader.import_deploy_checkpoints_config(job_src) as deploy_checkpoints_config:
+        assert len(deploy_checkpoints_config.checkpoint_details.checkpoints) == 2
+        assert (
+            deploy_checkpoints_config.checkpoint_details.base_model_id
+            == "unsloth/gemma-3-1b-it"
+        )
+        assert (
+            deploy_checkpoints_config.checkpoint_details.checkpoints[0].id
+            == "checkpoint-24"
+        )
+        assert (
+            deploy_checkpoints_config.checkpoint_details.checkpoints[1].id
+            == "checkpoint-42"
+        )
+
+
+def test_import_handles_training_project_with_deploy_checkpoints_config():
+    job_src = TEST_ROOT / "import" / "project_with_deploy_checkpoints_config.py"
+    with loader.import_training_project(job_src) as training_project:
+        assert training_project.name == "first-project"
+        assert training_project.job.compute.cpu_count == 4
+    with loader.import_deploy_checkpoints_config(job_src) as deploy_checkpoints_config:
+        assert len(deploy_checkpoints_config.checkpoint_details.checkpoints) == 2
+        assert (
+            deploy_checkpoints_config.checkpoint_details.base_model_id
+            == "unsloth/gemma-3-1b-it"
+        )
+        assert (
+            deploy_checkpoints_config.checkpoint_details.checkpoints[0].id
+            == "checkpoint-24"
+        )
+        assert (
+            deploy_checkpoints_config.checkpoint_details.checkpoints[1].id
+            == "checkpoint-42"
+        )

pyproject.toml

@@ -4,6 +4,8 @@
 
 from truss.base import custom_types, truss_config
 
+DEFAULT_LORA_RANK = 16
+
 
 class SecretReference(custom_types.SafeModel):
     name: str
@@ -24,6 +26,18 @@ def model_dump(self, *args, **kwargs):
             }
         return data
 
+    def to_truss_config(self) -> truss_config.Resources:
+        if self.accelerator:
+            return truss_config.Resources(
+                cpu=str(self.cpu_count),
+                memory=self.memory,
+                accelerator=self.accelerator,
+                node_count=self.node_count,
+            )
+        return truss_config.Resources(
+            cpu=str(self.cpu_count), memory=self.memory, node_count=self.node_count
+        )
+
 
 class CheckpointingConfig(custom_types.SafeModel):
     enabled: bool = False
@@ -57,3 +71,41 @@ class TrainingProject(custom_types.SafeModel):
     # TrainingProject is the wrapper around project config and job config. However, we exclude job
     # in serialization so just TrainingProject metadata is included in API requests.
     job: TrainingJob = pydantic.Field(exclude=True)
+
+
+class Checkpoint(custom_types.SafeModel):
+    training_job_id: str
+    id: str
+    name: str
+    lora_rank: Optional[int] = (
+        None  # lora rank will be fetched through the API if available.
+    )
+
+    def to_truss_config(self) -> truss_config.Checkpoint:
+        return truss_config.Checkpoint(
+            id=f"{self.training_job_id}/{self.id}", name=self.id
+        )
+
+
+class CheckpointDetails(custom_types.SafeModel):
+    download_folder: str = truss_config.DEFAULT_TRAINING_CHECKPOINT_FOLDER
+    base_model_id: Optional[str] = None
+    checkpoints: List[Checkpoint] = []
+
+    def to_truss_config(self) -> truss_config.CheckpointConfiguration:
+        checkpoints = [checkpoint.to_truss_config() for checkpoint in self.checkpoints]
+        return truss_config.CheckpointConfiguration(
+            checkpoints=checkpoints, download_folder=self.download_folder
+        )
+
+
+class DeployCheckpointsRuntime(custom_types.SafeModel):
+    environment_variables: Dict[str, Union[str, SecretReference]] = {}
+
+
+class DeployCheckpointsConfig(custom_types.SafeModel):
+    checkpoint_details: Optional[CheckpointDetails] = None
+    model_name: Optional[str] = None
+    deployment_name: Optional[str] = None
+    runtime: Optional[DeployCheckpointsRuntime] = None
+    compute: Optional[Compute] = None

truss-chains/truss_chains/deployment/deployment_client.py

@@ -2,13 +2,31 @@
 import importlib.util
 import os
 import pathlib
-from typing import Iterator
+from typing import Iterator, Type, TypeVar
 
 from truss_train import definitions
 
+T = TypeVar("T")
+
+
+@contextlib.contextmanager
+def import_training_project(
+    module_path: pathlib.Path,
+) -> Iterator[definitions.TrainingProject]:
+    with import_target(module_path, definitions.TrainingProject) as project:
+        yield project
+
+
+@contextlib.contextmanager
+def import_deploy_checkpoints_config(
+    module_path: pathlib.Path,
+) -> Iterator[definitions.DeployCheckpointsConfig]:
+    with import_target(module_path, definitions.DeployCheckpointsConfig) as config:
+        yield config
+
 
 @contextlib.contextmanager
-def import_target(module_path: pathlib.Path) -> Iterator[definitions.TrainingProject]:
+def import_target(module_path: pathlib.Path, target_type: Type[T]) -> Iterator[T]:
     module_name = module_path.stem
     if not os.path.isfile(module_path):
         raise ImportError(
@@ -24,13 +42,11 @@ def import_target(module_path: pathlib.Path) -> Iterator[definitions.TrainingPro
     spec.loader.exec_module(module)
 
     module_vars = (getattr(module, name) for name in dir(module))
-    training_projects = [
-        sym for sym in module_vars if isinstance(sym, definitions.TrainingProject)
-    ]
+    target = [sym for sym in module_vars if isinstance(sym, target_type)]
 
-    if len(training_projects) == 0:
-        raise ValueError(f"No `{definitions.TrainingProject}` was found.")
-    elif len(training_projects) > 1:
-        raise ValueError(f"Multiple `{definitions.TrainingProject}`s were found.")
+    if len(target) == 0:
+        raise ValueError(f"No `{target_type}` was found.")
+    elif len(target) > 1:
+        raise ValueError(f"Multiple `{target_type}`s were found.")
 
-    yield training_projects[0]
+    yield target[0]

truss-train/tests/import/deploy_checkpoints_config.py

@@ -1,15 +1,6 @@
 import pathlib
-from typing import Set
 
-SKLEARN = "sklearn"
-TENSORFLOW = "tensorflow"
-KERAS = "keras"
-XGBOOST = "xgboost"
-PYTORCH = "pytorch"
 CUSTOM = "custom"
-HUGGINGFACE_TRANSFORMER = "huggingface_transformer"
-LIGHTGBM = "lightgbm"
-
 
 _TRUSS_ROOT = pathlib.Path(__file__).parent.parent.resolve()
 
@@ -74,7 +65,6 @@
 TRUSS_DIR = "truss_dir"
 TRUSS_HASH = "truss_hash"
 
-HUGGINGFACE_TRANSFORMER_MODULE_NAME: Set[str] = set({})
 
 INFERENCE_SERVER_PORT = 8080
 
@@ -84,7 +74,7 @@
 
 TRTLLM_SPEC_DEC_TARGET_MODEL_NAME = "target"
 TRTLLM_SPEC_DEC_DRAFT_MODEL_NAME = "draft"
-TRTLLM_BASE_IMAGE = "baseten/briton-server:v0.18.1-cd81637"
+TRTLLM_BASE_IMAGE = "baseten/briton-server:v0.18.1-1856fb5"
 # TODO: build the image so that the default path `python3` can be used - then remove here.
 TRTLLM_PYTHON_EXECUTABLE = "/usr/local/briton/venv/bin/python"
 BEI_TRTLLM_BASE_IMAGE = "baseten/bei:0.0.23"