truss 0.11.3__py3-none-any.whl → 0.11.4__py3-none-any.whl

truss/remote/baseten/api.py

@@ -6,6 +6,7 @@ import requests
 
 from truss.remote.baseten import custom_types as b10_types
 from truss.remote.baseten.auth import ApiKey, AuthService
+from truss.remote.baseten.custom_types import APIKeyCategory
 from truss.remote.baseten.error import ApiError
 from truss.remote.baseten.rest_client import RestAPIClient
 from truss.remote.baseten.utils.transfer import base64_encoded_json_str
@@ -560,6 +561,13 @@ class BasetenApi:
     def get_all_secrets(self) -> Any:
         return self._rest_api_client.get("v1/secrets")
 
+    # NOTE(Tyron): `name` is required because all official
+    # Baseten API keys should have a descriptive name.
+    def create_api_key(self, api_key_type: APIKeyCategory, name: str) -> Any:
+        return self._rest_api_client.post(
+            "v1/api_keys", body={"type": api_key_type.value, "name": name}
+        )
+
     def upsert_training_project(self, training_project):
         resp_json = self._rest_api_client.post(
             "v1/training_projects",

truss/remote/baseten/custom_types.py

@@ -155,3 +155,10 @@ class GetCacheSummaryResponseV1(pydantic.BaseModel):
     file_summaries: list[FileSummary] = pydantic.Field(
         description="List of files in the cache"
     )
+
+
+class APIKeyCategory(Enum):
+    PERSONAL = "PERSONAL"
+    WORKSPACE_MANAGE_ALL = "WORKSPACE_MANAGE_ALL"
+    WORKSPACE_EXPORT_METRICS = "WORKSPACE_EXPORT_METRICS"
+    WORKSPACE_INVOKE = "WORKSPACE_INVOKE"

truss/templates/control/requirements.txt

@@ -7,6 +7,7 @@ python-json-logger>=2.0.2
 tenacity>=8.1.0
  # To avoid divergence, this should follow the latest release.
 truss==0.11.1
-uvicorn>=0.24.0
+# NB(nikhil): Uvicorn 0.36.0 has breaking changes for the event loop, so we pin to a lower version.
+uvicorn>=0.24.0,<0.36.0
 uvloop>=0.19.0
 websockets>=10.0

truss/templates/server/requirements.txt

@@ -19,6 +19,7 @@ python-json-logger>=2.0.2
 pyyaml>=6.0.0
 requests>=2.31.0
 truss-transfer==0.0.30
-uvicorn>=0.24.0
+# NB(nikhil): Uvicorn 0.36.0 has breaking changes for the event loop, so we pin to a lower version.
+uvicorn>=0.24.0,<0.36.0
 uvloop>=0.19.0
 websockets>=10.0

{truss-0.11.3.dist-info → truss-0.11.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: truss
-Version: 0.11.3
+Version: 0.11.4
 Summary: A seamless bridge from model development to model delivery
 Project-URL: Repository, https://github.com/basetenlabs/truss
 Project-URL: Homepage, https://truss.baseten.co

{truss-0.11.3.dist-info → truss-0.11.4.dist-info}/RECORD

@@ -52,10 +52,10 @@ truss/patch/truss_dir_patch_applier.py,sha256=ALnaVnu96g0kF2UmGuBFTua3lrXpwAy4sG
 truss/remote/remote_factory.py,sha256=-0gLh_yIyNDgD48Q6sR8Yo5dOMQg84lrHRvn_XR0n4s,3585
 truss/remote/truss_remote.py,sha256=TEe6h6by5-JLy7PMFsDN2QxIY5FmdIYN3bKvHHl02xM,8440
 truss/remote/baseten/__init__.py,sha256=XNqJW1zyp143XQc6-7XVwsUA_Q_ZJv_ausn1_Ohtw9Y,176
-truss/remote/baseten/api.py,sha256=lJOt2i3tu0ZeCh4B_-hpfpjcZKgTHVnkxraooK7TUHw,24699
+truss/remote/baseten/api.py,sha256=5XsB0mSfFD2YGybco4aPGMQalSvhsqv-QgzW4A66r7k,25081
 truss/remote/baseten/auth.py,sha256=tI7s6cI2EZgzpMIzrdbILHyGwiHDnmoKf_JBhJXT55E,776
 truss/remote/baseten/core.py,sha256=uxtmBI9RAVHu1glIEJb5Q4ccJYLeZM1Cp5Svb9W68Yw,21965
-truss/remote/baseten/custom_types.py,sha256=1OfbZwT-n7mrYQ4ygfWuvOnENedZ4L6zOSFMAhBAVqI,4509
+truss/remote/baseten/custom_types.py,sha256=bYrfTzGgYr6FDoya0omyadCLSTcTc-83U2scQORyUj0,4715
 truss/remote/baseten/error.py,sha256=3TNTwwPqZnr4NRd9Sl6SfLUQR2fz9l6akDPpOntTpzA,578
 truss/remote/baseten/remote.py,sha256=Se8AES5mk8jxa8S9fN2DSG7wnsaV7ftRjJ4Uwc_w_S0,22544
 truss/remote/baseten/rest_client.py,sha256=_t3CWsWARt2u0C0fDsF4rtvkkHe-lH7KXoPxWXAkKd4,1185
@@ -72,7 +72,7 @@ truss/templates/cache_requirements.txt,sha256=xoPoJ-OVnf1z6oq_RVM3vCr3ionByyqMLj
 truss/templates/copy_cache_files.Dockerfile.jinja,sha256=Os5zFdYLZ_AfCRGq4RcpVTObOTwL7zvmwYcvOzd_Zqo,126
 truss/templates/docker_server_requirements.txt,sha256=PyhOPKAmKW1N2vLvTfLMwsEtuGpoRrbWuNo7tT6v2Mc,18
 truss/templates/server.Dockerfile.jinja,sha256=CUYnF_hgxPGq2re7__0UPWlwzOHMoFkxp6NVKi3U16s,7071
-truss/templates/control/requirements.txt,sha256=tJGr83WoE0CZm2FrloZ9VScK84q-_FTuVXjDYrexhW0,250
+truss/templates/control/requirements.txt,sha256=nqqNmlTwFeV8sV4fqwItwzzd_egADBP_e-cEopXBJ4k,358
 truss/templates/control/control/application.py,sha256=jYeta6hWe1SkfLL3W4IDmdYjg3ZuKqI_UagWYs5RB_E,3793
 truss/templates/control/control/endpoints.py,sha256=VQ1lvZjFvR091yRkiFdvXw1Q7PiNGXT9rJwY7_sX6yg,11828
 truss/templates/control/control/server.py,sha256=R4Y219i1dcz0kkksN8obLoX-YXWGo9iW1igindyG50c,3128
@@ -96,7 +96,7 @@ truss/templates/docker_server/supervisord.conf.jinja,sha256=dd37fwZE--cutrvOUCqE
 truss/templates/server/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 truss/templates/server/main.py,sha256=kWXrdD8z8IpamyWxc8qcvd5ck9gM1Kz2QH5qHJCnmOQ,222
 truss/templates/server/model_wrapper.py,sha256=k75VVISwwlsx5EGb82UZsu8kCM_i6Yi3-Hd0-Kpm1yo,42055
-truss/templates/server/requirements.txt,sha256=XblmpfxAmRo3X1V_9oMj8yjdpZ5Wk-C2oa3z6nq4OGw,672
+truss/templates/server/requirements.txt,sha256=fMLT-6FuY-j3dyJdOosR5fKqcBOkCZOob7Pl_Y0feg8,780
 truss/templates/server/truss_server.py,sha256=noXfGJMsKIhgF4oI_8LC1UHkcx8Vg8nGSITZJ_bkRFQ,19598
 truss/templates/server/common/__init__.py,sha256=qHIqr68L5Tn4mV6S-PbORpcuJ4jmtBR8aCuRTIWDvNo,85
 truss/templates/server/common/errors.py,sha256=My0P6-Y7imVTICIhazHT0vlSu3XJDH7As06OyVzu4Do,8589
@@ -348,14 +348,14 @@ truss_chains/__init__.py,sha256=QDw1YwdqMaQpz5Oltu2Eq2vzEX9fDrMoqnhtbeh60i4,1278
 truss_chains/framework.py,sha256=CS7tSegPe2Q8UUT6CDkrtSrB3utr_1QN1jTEPjrj5Ug,67519
 truss_chains/private_types.py,sha256=6CaQEPawFLXjEbJ-01lqfexJtUIekF_q61LNENWegFo,8917
 truss_chains/public_api.py,sha256=0AXV6UdZIFAMycUNG_klgo4aLFmBZeKGfrulZEWzR0M,9532
-truss_chains/public_types.py,sha256=Am1vc5pWuEDs65UQp8Be4iOU05kMXjPG3QrJfNXDmHs,29225
+truss_chains/public_types.py,sha256=RPr8jgKO_F_26F7H3CpwbidL-6euoKPdFHVpEIpYqrQ,29415
 truss_chains/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 truss_chains/pydantic_numpy.py,sha256=MG8Ji_Inwo_JSfM2n7TPj8B-nbrBlDYsY3SOeBwD8fE,4289
 truss_chains/streaming.py,sha256=DGl2LEAN67YwP7Nn9MK488KmYc4KopWmcHuE6WjyO1Q,12521
 truss_chains/utils.py,sha256=LvpCG2lnN6dqPqyX3PwLH9tyjUzqQN3N4WeEFROMHak,6291
 truss_chains/deployment/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-truss_chains/deployment/code_gen.py,sha256=IBOYdhsWUyW_sBVhlEQAhvwxKcsGflDjgmR-1HyJJLg,32666
-truss_chains/deployment/deployment_client.py,sha256=2paNyBjrpFTxROP0YrmJMUlH6o8mrkF-iPms7VJhLdA,34017
+truss_chains/deployment/code_gen.py,sha256=397FiSNZuW59J3Ma7N9GKGfvG_87BNFAXCIV8BW41t0,32669
+truss_chains/deployment/deployment_client.py,sha256=OoqkO3daktYzR2YsIcDvsuGfjR05X2K7QlA7wvFduzc,34208
 truss_chains/reference_code/reference_chainlet.py,sha256=5feSeqGtrHDbldkfZCfX2R5YbbW0Uhc35mhaP2pXrHw,1340
 truss_chains/reference_code/reference_model.py,sha256=emH3hb23E_nbP98I37PGp1Xk1hz3g3lQ00tiLo55cSM,322
 truss_chains/remote_chainlet/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -368,8 +368,8 @@ truss_train/deployment.py,sha256=lWWANSuzBWu2M4oK4qD7n-oVR1JKdmw2Pn5BJQHg-Ck,307
 truss_train/loader.py,sha256=0o66EjBaHc2YY4syxxHVR4ordJWs13lNXnKjKq2wq0U,1630
 truss_train/public_api.py,sha256=9N_NstiUlmBuLUwH_fNG_1x7OhGCytZLNvqKXBlStrM,1220
 truss_train/restore_from_checkpoint.py,sha256=KmJuTUVpvtvlkEClcmllxAF2TKgbp-FuzfblfGh06XA,1239
-truss-0.11.3.dist-info/METADATA,sha256=tzcaX0tJ3oM9TfBC0FGlbh3No2PfcyBO4Yfg3IPMkKE,6669
-truss-0.11.3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-truss-0.11.3.dist-info/entry_points.txt,sha256=-MwKfHHQHQ6j0HqIgvxrz3CehCmczDLTD-OsRHnjjuU,130
-truss-0.11.3.dist-info/licenses/LICENSE,sha256=FTqGzu85i-uw1Gi8E_o0oD60bH9yQ_XIGtZbA1QUYiw,1064
-truss-0.11.3.dist-info/RECORD,,
+truss-0.11.4.dist-info/METADATA,sha256=sDn320OamuGyMaRn5BvlChKqzViyvON94QTEcPCgBC4,6669
+truss-0.11.4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+truss-0.11.4.dist-info/entry_points.txt,sha256=-MwKfHHQHQ6j0HqIgvxrz3CehCmczDLTD-OsRHnjjuU,130
+truss-0.11.4.dist-info/licenses/LICENSE,sha256=FTqGzu85i-uw1Gi8E_o0oD60bH9yQ_XIGtZbA1QUYiw,1064
+truss-0.11.4.dist-info/RECORD,,

truss_chains/deployment/code_gen.py

@@ -764,13 +764,13 @@ def _gen_truss_config(
             config.external_package_dirs.append(ext_dir.abs_path)
     config.use_local_src = use_local_src
 
-    if public_types._BASETEN_API_SECRET_NAME not in config.secrets:
-        config.secrets[public_types._BASETEN_API_SECRET_NAME] = (
+    if public_types.CHAIN_API_KEY_SECRET_NAME not in config.secrets:
+        config.secrets[public_types.CHAIN_API_KEY_SECRET_NAME] = (
             public_types.SECRET_DUMMY
         )
     else:
         logging.info(
-            f"Chains automatically add {public_types._BASETEN_API_SECRET_NAME} "
+            f"Chains automatically add {public_types.CHAIN_API_KEY_SECRET_NAME} "
             "to secrets - no need to manually add it."
         )
     config.model_cache = truss_config.ModelCache(assets.cached)

truss_chains/deployment/deployment_client.py

@@ -303,7 +303,7 @@ def _push_service_docker(
     options: private_types.PushOptionsLocalDocker,
 ) -> str:
     th = truss_handle.TrussHandle(truss_dir)
-    th.add_secret(public_types._BASETEN_API_SECRET_NAME, options.baseten_chain_api_key)
+    th.add_secret(public_types.CHAIN_API_KEY_SECRET_NAME, options.baseten_chain_api_key)
     container = th.docker_run(
         local_port=None,
         detach=True,
@@ -536,14 +536,21 @@ def _create_baseten_chain(
 def _create_chains_secret_if_missing(remote_provider: b10_remote.BasetenRemote) -> None:
     secrets_info = remote_provider.api.get_all_secrets()
     secret_names = {sec["name"] for sec in secrets_info["secrets"]}
-    if public_types._BASETEN_API_SECRET_NAME not in secret_names:
+
+    if public_types.CHAIN_API_KEY_SECRET_NAME not in secret_names:
         logging.info(
             "It seems you are using chains for the first time, since there "
-            f"is no `{public_types._BASETEN_API_SECRET_NAME}` secret on baseten. "
+            f"is no `{public_types.CHAIN_API_KEY_SECRET_NAME}` secret on baseten. "
             "Creating secret automatically."
         )
+
+        workspace_api_key = remote_provider.api.create_api_key(
+            api_key_type=b10_types.APIKeyCategory.WORKSPACE_INVOKE,
+            name=public_types.CHAIN_API_KEY_NAME,
+        )["api_key"]
+
         remote_provider.api.upsert_secret(
-            public_types._BASETEN_API_SECRET_NAME, remote_provider.api.auth_token.value
+            public_types.CHAIN_API_KEY_SECRET_NAME, workspace_api_key
         )
 
 

truss_chains/public_types.py

@@ -28,7 +28,12 @@ DEFAULT_CONCURRENCY_LIMIT = 300
 CpuCountT = Literal["cpu_count"]
 CPU_COUNT: CpuCountT = "cpu_count"
 
-_BASETEN_API_SECRET_NAME = "baseten_chain_api_key"
+
+# NOTE(Tyron): This is a secret that points to an API key value.
+CHAIN_API_KEY_SECRET_NAME = "baseten_chain_api_key"
+
+# NOTE(Tyron): This is the actual API key pointed to by the above secret.
+CHAIN_API_KEY_NAME = "baseten-chain-api-key"
 
 _K = TypeVar("_K", contravariant=True)
 _V = TypeVar("_V", covariant=True)
@@ -738,14 +743,14 @@ class DeploymentContext(custom_types.SafeModelNonSerializable):
             )
         error_msg = (
             "For using chains, it is required to setup a an API key with name "
-            f"`{_BASETEN_API_SECRET_NAME}` on Baseten to allow chain Chainlet to "
+            f"`{CHAIN_API_KEY_SECRET_NAME}` on Baseten to allow chain Chainlet to "
             "call other Chainlets. For local execution, secrets can be provided "
             "to `run_local`."
         )
-        if _BASETEN_API_SECRET_NAME not in self.secrets:
+        if CHAIN_API_KEY_SECRET_NAME not in self.secrets:
             raise MissingDependencyError(error_msg)
 
-        api_key = self.secrets[_BASETEN_API_SECRET_NAME]
+        api_key = self.secrets[CHAIN_API_KEY_SECRET_NAME]
         if api_key == SECRET_DUMMY:
             raise MissingDependencyError(
                 f"{error_msg}. Retrieved dummy value of `{api_key}`."