Add s3 session token support (#718)

* Add s3 session token support

* Add more documentation. Re-order cache Dockerfile for better caching

* Update pyproject.toml

Author: bolasim

docs/guides/model-cache.mdx

@@ -135,6 +135,7 @@ model_cache:
 If you are accessing a public GCS bucket, you can ignore the subsequent steps, but make sure you set an appropriate appropriate policy on your bucket. Users should be able to list and view all files. Otherwise, the model build will fail.
 
 However, for a private S3 bucket, you need to first find your `aws_access_key_id`, `aws_secret_access_key`, and `aws_region` in your AWS dashboard. Create a file named `s3_credentials.json`. Inside this file, add the credentials that you identified earlier as shown below. Place this file into the `data` directory of your Truss.
+The key `aws_session_token` can be included, but is optional.
 
 Here is an example of how your `s3_credentials.json` file should look:
 
@@ -156,6 +157,32 @@ your-truss
 |. └── s3_credentials.json
 ```
 
+When you are generating credentials, make sure that the resulting keys have at minimum the following IAM policy:
+
+```json
+{
+        "Version": "2012-10-17",
+        "Statement": [
+            {
+                "Action": [
+                    "s3:GetObject",
+                    "s3:ListObjects",
+                ],
+                "Effect": "Allow",
+                "Resource": ["arn:aws:s3:::S3_BUCKET/PATH_TO_MODEL/*"]
+            },
+            {
+                "Action": [
+                    "s3:ListBucket",
+                ],
+                "Effect": "Allow",
+                "Resource": ["arn:aws:s3:::S3_BUCKET"]
+            }
+        ]
+    }
+```
+
+
 <Warning>
 If you are using version control, like git, for your Truss, make sure to add `s3_credentials.json` to your `.gitignore` file. You don't want to accidentally expose your service account key.
 </Warning>

pyproject.toml

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

truss/contexts/image_builder/cache_warmer.py

@@ -47,6 +47,7 @@ class AWSCredentials:
     access_key_id: str
     secret_access_key: str
     region: str
+    session_token: Optional[str]
 
 
 def parse_s3_credentials_file(key_file_path: str) -> AWSCredentials:
@@ -67,6 +68,7 @@ def parse_s3_credentials_file(key_file_path: str) -> AWSCredentials:
         access_key_id=data["aws_access_key_id"],
         secret_access_key=data["aws_secret_access_key"],
         region=data["aws_region"],
+        session_token=data.get("aws_session_token", None),
     )
 
     return aws_sa
@@ -179,6 +181,7 @@ def download_to_cache(self):
                 aws_access_key_id=s3_credentials.access_key_id,
                 aws_secret_access_key=s3_credentials.secret_access_key,
                 region_name=s3_credentials.region,
+                aws_session_token=s3_credentials.session_token,
                 config=Config(signature_version="s3v4"),
             )
         else:

truss/contexts/image_builder/serving_image_builder.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import json
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from pathlib import Path
@@ -27,6 +26,10 @@
     TEMPLATES_DIR,
     TRITON_SERVER_CODE_DIR,
 )
+from truss.contexts.image_builder.cache_warmer import (
+    AWSCredentials,
+    parse_s3_credentials_file,
+)
 from truss.contexts.image_builder.image_builder import ImageBuilder
 from truss.contexts.image_builder.util import (
     TRUSS_BASE_IMAGE_VERSION_TAG,
@@ -136,10 +139,15 @@ def list_files(self, revision=None):
         s3_credentials_file = self.data_dir / S3_CREDENTIALS
 
         if s3_credentials_file.exists():
-            AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY = parse_s3_service_account_file(
+            s3_credentials: AWSCredentials = parse_s3_credentials_file(
                 self.data_dir / S3_CREDENTIALS
             )
-            session = boto3.Session(AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY)
+            session = boto3.Session(
+                aws_access_key_id=s3_credentials.access_key_id,
+                aws_secret_access_key=s3_credentials.secret_access_key,
+                aws_session_token=s3_credentials.session_token,
+                region_name=s3_credentials.region,
+            )
             s3 = session.resource("s3")
         else:
             s3 = boto3.resource("s3", config=Config(signature_version=UNSIGNED))
@@ -274,22 +282,6 @@ class CachedFile:
     dst: str
 
 
-def parse_s3_service_account_file(file_path):
-    # open the json file
-    with open(file_path, "r") as f:
-        data = json.load(f)
-
-    # validate the data
-    if "aws_access_key_id" not in data or "aws_secret_access_key" not in data:
-        raise ValueError("Invalid AWS credentials file")
-
-    # parse the data
-    aws_access_key_id = data["aws_access_key_id"]
-    aws_secret_access_key = data["aws_secret_access_key"]
-
-    return aws_access_key_id, aws_secret_access_key
-
-
 def update_model_key(config: TrussConfig) -> str:
     server_name = config.build.model_server
 

truss/templates/cache.Dockerfile.jinja

@@ -7,15 +7,16 @@ WORKDIR /app
 ENV HUGGING_FACE_HUB_TOKEN {{hf_access_token}}
 {% endif %}
 
-{% for credential in credentials_to_cache  %}
-COPY ./{{credential}} /app/{{credential}}
-{% endfor %}
-
 RUN apt-get -y update; apt-get -y install curl; curl -s https://baseten-public.s3.us-west-2.amazonaws.com/bin/b10cp-5fe8dc7da-linux-amd64 -o /app/b10cp; chmod +x /app/b10cp
 ENV B10CP_PATH_TRUSS /app/b10cp
 COPY ./cache_requirements.txt /app/cache_requirements.txt
 RUN pip install -r /app/cache_requirements.txt --no-cache-dir && rm -rf /root/.cache/pip
 COPY ./cache_warmer.py /cache_warmer.py
+
+{% for credential in credentials_to_cache  %}
+COPY ./{{credential}} /app/{{credential}}
+{% endfor %}
+
 {% for repo, hf_dir in models.items() %}
         {% for file in hf_dir.files %}
 {{ "RUN --mount=type=secret,id=" + hf_access_token_file_name + ",dst=/etc/secrets/" + hf_access_token_file_name if use_hf_secret else "RUN" }} python3 /cache_warmer.py {{file}} {{repo}} {% if hf_dir.revision != None %}{{hf_dir.revision}}{% endif %}