Merge pull request #570 from basetenlabs/main

Release 0.6.0

Author: squidarth

docs/mint.json

@@ -179,5 +179,10 @@
       "name": "Issues",
       "url": "https://github.com/basetenlabs/truss/issues"
     }
-  ]
+  ],
+  "analytics": {
+    "gtm": {
+        "tagId": "GTM-WXD4NQTW"
+    }
+  }
 }

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "truss"
-version = "0.5.9rc1"
+version = "0.6.0"
 description = "A seamless bridge from model development to model delivery"
 license = "MIT"
 readme = "README.md"

truss/build.py

@@ -218,6 +218,7 @@ def init(
     bundled_packages: Optional[List[str]] = None,
     trainable: bool = False,
     build_config: Optional[Build] = None,
+    model_name: Optional[str] = None,
 ) -> TrussHandle:
     """
     Initialize an empty placeholder Truss. A Truss is a build context designed
@@ -230,6 +231,7 @@ def init(
                           Truss in. The directory is created if it doesn't exist.
     """
     config = TrussConfig(
+        model_name=model_name,
         python_version=map_to_supported_python_version(infer_python_version()),
     )
 

truss/cli/cli.py

@@ -11,9 +11,11 @@
 import rich
 import rich_click as click
 import truss
-from truss.cli.create import select_server_backend
+from InquirerPy import inquirer
+from truss.cli.create import ask_name, select_server_backend
 from truss.remote.remote_cli import inquire_model_name, inquire_remote_name
 from truss.remote.remote_factory import RemoteFactory
+from truss.truss_config import ModelServer
 
 logging.basicConfig(level=logging.INFO)
 
@@ -51,6 +53,8 @@ def error_handling(f: Callable[..., object]):
     def wrapper(*args, **kwargs):
         try:
             f(*args, **kwargs)
+        except click.UsageError as e:
+            raise e  # You can re-raise the exception or handle it different
         except Exception as e:
             click.echo(e)
 
@@ -93,8 +97,15 @@ def image():
     default=False,
     help="Create a trainable truss.",
 )
+@click.option(
+    "-b",
+    "--backend",
+    show_default=True,
+    default=ModelServer.TrussServer.value,
+    type=click.Choice([server.value for server in ModelServer]),
+)
 @error_handling
-def init(target_directory, trainable) -> None:
+def init(target_directory, trainable, backend) -> None:
     """Create a new truss.
 
     TARGET_DIRECTORY: A Truss is created in this directory
@@ -104,13 +115,15 @@ def init(target_directory, trainable) -> None:
             f'Error: Directory "{target_directory}" already exists and cannot be overwritten.'
         )
     tr_path = Path(target_directory)
-    build_config = select_server_backend()
+    build_config = select_server_backend(ModelServer[backend])
+    model_name = ask_name()
     truss.init(
         target_directory=target_directory,
         trainable=trainable,
         build_config=build_config,
+        model_name=model_name,
     )
-    click.echo(f"Truss was created in {tr_path}")
+    click.echo(f"Truss {model_name} was created in {tr_path.absolute()}")
 
 
 @image.command()  # type: ignore
@@ -184,10 +197,18 @@ def run(target_directory: str, build_dir: Path, tag, port, attach) -> None:
     required=False,
     help="Name of the remote in .trussrc to patch changes to",
 )
+@click.option(
+    "--logs",
+    is_flag=True,
+    show_default=True,
+    default=False,
+    help="Automatically open remote logs tab",
+)
 @error_handling
 def watch(
     target_directory: str,
     remote: str,
+    logs: bool,
 ) -> None:
     """
     Seamless remote development with truss
@@ -211,9 +232,13 @@ def watch(
         sys.exit(1)
 
     logs_url = remote_provider.get_remote_logs_url(model_name)  # type: ignore[attr-defined]
-    rich.print(f"🪵 View logs for your deployment at {logs_url}")
-    webbrowser.open(logs_url)
-    rich.print(f"👀 Watching for changes to truss at '{target_directory}' ...")
+    rich.print(f"🪵  View logs for your deployment at {logs_url}")
+    if not logs:
+        logs = inquirer.confirm(
+            message="🗂  Open logs in a new tab?", default=True
+        ).execute()
+    if logs:
+        webbrowser.open_new_tab(logs_url)
     remote_provider.sync_truss_to_dev_version_by_name(model_name, target_directory)  # type: ignore
 
 
@@ -273,14 +298,18 @@ def predict(
 
     model_name = tr.spec.config.model_name
     if not model_name:
-        raise ValueError("Model name not set. Did you `truss push`?")
+        raise click.UsageError(
+            "You must provide exactly one of '--data (-d)' or '--file (-f)' options."
+        )
 
     if data is not None:
         request_data = json.loads(data)
     elif file is not None:
         request_data = json.loads(Path(file).read_text())
     else:
-        raise ValueError("At least one of request or request-file must be supplied.")
+        raise click.UsageError(
+            "You must provide exactly one of '--data (-d)' or '--file (-f)' options."
+        )
 
     service = remote_provider.get_baseten_service(model_name, published)  # type: ignore
     result = service.predict(request_data)

truss/cli/create.py

@@ -24,14 +24,7 @@
 }
 
 
-def select_server_backend() -> Build:
-    server_backend = ModelServer[
-        inquirer.select(
-            message="Select a server:",
-            choices=[s.value for s in ModelServer],
-            default=ModelServer.TrussServer.value,
-        ).execute()
-    ]
+def select_server_backend(server_backend: ModelServer) -> Build:
     follow_up_questions = REQUIRED_ARGUMENTS.get(server_backend)
     args = {}
     if follow_up_questions:
@@ -41,3 +34,7 @@ def select_server_backend() -> Build:
             ).execute()
 
     return Build(model_server=server_backend, arguments=args)
+
+
+def ask_name() -> str:
+    return inquirer.text(message="What's the name of your model?").execute()

truss/contexts/local_loader/truss_file_syncer.py

@@ -2,6 +2,8 @@
 from pathlib import Path
 from threading import Thread
 
+import rich
+
 
 class TrussFilesSyncer(Thread):
     """Daemon thread that watches for changes in the user's Truss and syncs to running service."""
@@ -29,6 +31,10 @@ def run(self) -> None:
         # disable watchfiles logger
         logging.getLogger("watchfiles.main").disabled = True
 
+        rich.print(f"🚰 Attempting to sync truss at '{self.watch_path}' with remote")
+        self.remote.patch(self.watch_path, self._logger)
+
+        rich.print(f"👀 Watching for changes to truss at '{self.watch_path}' ...")
         for _ in watch(
             self.watch_path, watch_filter=self.watch_filter, raise_interrupt=False
         ):

truss/remote/baseten/remote.py

@@ -2,6 +2,7 @@
 from pathlib import Path
 
 import yaml
+from requests import ReadTimeout
 from truss.contexts.local_loader.truss_file_syncer import TrussFilesSyncer
 from truss.local.local_config_handler import LocalConfigHandler
 from truss.remote.baseten.api import BasetenApi
@@ -137,6 +138,11 @@ def patch(
         except yaml.parser.ParserError:
             logger.error("Unable to parse config file")
             return
+        except ValueError:
+            logger.error(
+                f"Error when reading truss from directory {watch_path}", exc_info=True
+            )
+            return
         model_name = truss_handle.spec.config.model_name
         dev_version = get_dev_version_info(self._api, model_name)  # type: ignore
         truss_hash = dev_version.get("truss_hash", None)
@@ -145,7 +151,7 @@ def patch(
         try:
             patch_request = truss_handle.calc_patch(truss_hash)
         except Exception:
-            logger.error("Failed to calculate patch")
+            logger.error("Failed to calculate patch, bailing on patching")
             return
         if patch_request:
             if (
@@ -154,7 +160,16 @@ def patch(
             ):
                 logger.info("No changes observed, skipping deployment")
                 return
-            resp = self._api.patch_draft_truss(model_name, patch_request)
+            try:
+                resp = self._api.patch_draft_truss(model_name, patch_request)
+            except ReadTimeout:
+                logger.error(
+                    "Read Timeout when attempting to connect to remote. Bailing on patching"
+                )
+                return
+            except Exception:
+                logger.error("Failed to patch draft deployment, bailing on patching")
+                return
             if not resp["succeeded"]:
                 needs_full_deploy = resp.get("needs_full_deploy", None)
                 if needs_full_deploy:
@@ -169,6 +184,6 @@ def patch(
                 logger.info(
                     resp.get(
                         "success_message",
-                        f"Model {model_name} patched successfully.",
+                        f"Model {model_name} patched successfully",
                     )
                 )

truss/remote/remote_cli.py

@@ -9,13 +9,11 @@
 def inquire_remote_config() -> RemoteConfig:
     # TODO(bola): extract questions from remote
     rich.print("💻 Let's add a Baseten remote!")
-    remote_url = inquirer.text(
-        message="🌐 Baseten remote url:",
-        default="https://app.baseten.co",
-        qmark="",
-    ).execute()
+    # If users need to adjust the remote url, they
+    # can do so manually in the .trussrc file.
+    remote_url = "https://app.baseten.co"
     api_key = inquirer.secret(
-        message="🤫 Quiety paste your API_KEY:",
+        message="🤫 Quietly paste your API_KEY:",
         qmark="",
     ).execute()
     return RemoteConfig(

truss/templates/custom/model/model.py

@@ -1,19 +1,32 @@
-from typing import Any
+"""
+The `Model` class is an interface between the ML model that you're packaging and the model
+server that you're running it on.
+
+The main methods to implement here are:
+* `load`: runs exactly once when the model server is spun up or patched and loads the
+   model onto the model server. Include any logic for initializing your model, such
+   as downloading model weights and loading the model into memory.
+* `predict`: runs every time the model server is called. Include any logic for model
+  inference and return the model output.
+
+See https://truss.baseten.co/quickstart for more.
+"""
 
 
 class Model:
-    def __init__(self, **kwargs) -> None:
-        self._data_dir = kwargs["data_dir"]
-        self._config = kwargs["config"]
-        self._secrets = kwargs["secrets"]
+    def __init__(self, **kwargs):
+        # Uncomment the following to get access
+        # to various parts of the Truss config.
+
+        # self._data_dir = kwargs["data_dir"]
+        # self._config = kwargs["config"]
+        # self._secrets = kwargs["secrets"]
         self._model = None
 
     def load(self):
         # Load model here and assign to self._model.
         pass
 
-    def predict(self, model_input: Any) -> Any:
-        model_output = {}
-        # Invoke model on model_input and calculate predictions here.
-        model_output["predictions"] = []
-        return model_output
+    def predict(self, model_input):
+        # Run model inference here
+        return model_input

truss/tests/test_config.py

@@ -129,8 +129,8 @@ def generate_default_config():
         "requirements": [],
         "resources": {
             "accelerator": None,
-            "cpu": "500m",
-            "memory": "512Mi",
+            "cpu": "1",
+            "memory": "2Gi",
             "use_gpu": False,
         },
         "secrets": {},
@@ -153,8 +153,8 @@ def test_default_config_not_crowded_end_to_end():
 requirements: []
 resources:
   accelerator: null
-  cpu: 500m
-  memory: 512Mi
+  cpu: '1'
+  memory: 2Gi
   use_gpu: false
 secrets: {}
 system_packages: []
@@ -192,8 +192,8 @@ def test_non_default_train():
     updated_train = {
         "resources": {
             "accelerator": "A10G",
-            "cpu": "500m",
-            "memory": "512Mi",
+            "cpu": "1",
+            "memory": "2Gi",
             "use_gpu": True,
         },
         "variables": {},

truss/tests/test_truss_handle.py

@@ -856,8 +856,8 @@ def generate_default_config():
         "requirements": [],
         "resources": {
             "accelerator": None,
-            "cpu": "500m",
-            "memory": "512Mi",
+            "cpu": "1",
+            "memory": "2Gi",
             "use_gpu": False,
         },
         "secrets": {},

truss/truss_config.py

@@ -29,8 +29,8 @@
 DEFAULT_SPEC_VERSION = "2.0"
 DEFAULT_PREDICT_CONCURRENCY = 1
 
-DEFAULT_CPU = "500m"
-DEFAULT_MEMORY = "512Mi"
+DEFAULT_CPU = "1"
+DEFAULT_MEMORY = "2Gi"
 DEFAULT_USE_GPU = False
 
 DEFAULT_TRAINING_CLASS_FILENAME = "train.py"