Change default truss.

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"