45
45
CONFIG_FILE = "config.yaml"
46
46
47
47
HF_ACCESS_TOKEN_SECRET_NAME = "hf_access_token"
48
+ HF_ACCESS_TOKEN_FILE_NAME = "hf-access-token"
48
49
49
50
50
51
def create_triton_build_dir (config : TrussConfig , build_dir : Path , truss_dir : Path ):
@@ -234,7 +235,9 @@ def update_config_and_gather_files(
234
235
return get_files_to_cache (config , truss_dir , build_dir )
235
236
236
237
237
- def create_tgi_build_dir (config : TrussConfig , build_dir : Path , truss_dir : Path ):
238
+ def create_tgi_build_dir (
239
+ config : TrussConfig , build_dir : Path , truss_dir : Path , use_hf_secret : bool
240
+ ):
238
241
copy_tree_path (truss_dir , build_dir )
239
242
240
243
if not build_dir .exists ():
@@ -258,6 +261,8 @@ def create_tgi_build_dir(config: TrussConfig, build_dir: Path, truss_dir: Path):
258
261
data_dir_exists = data_dir .exists (),
259
262
credentials_exists = credentials_file .exists (),
260
263
cached_files = cached_file_paths ,
264
+ use_hf_secret = use_hf_secret ,
265
+ hf_access_token_file_name = HF_ACCESS_TOKEN_FILE_NAME ,
261
266
)
262
267
dockerfile_filepath = build_dir / "Dockerfile"
263
268
dockerfile_filepath .write_text (dockerfile_content )
@@ -279,7 +284,9 @@ def create_tgi_build_dir(config: TrussConfig, build_dir: Path, truss_dir: Path):
279
284
supervisord_filepath .write_text (supervisord_contents )
280
285
281
286
282
- def create_vllm_build_dir (config : TrussConfig , build_dir : Path , truss_dir : Path ):
287
+ def create_vllm_build_dir (
288
+ config : TrussConfig , build_dir : Path , truss_dir : Path , use_hf_secret
289
+ ):
283
290
copy_tree_path (truss_dir , build_dir )
284
291
285
292
server_endpoint_config = {
@@ -312,6 +319,8 @@ def create_vllm_build_dir(config: TrussConfig, build_dir: Path, truss_dir: Path)
312
319
data_dir_exists = data_dir .exists (),
313
320
credentials_exists = credentials_file .exists (),
314
321
cached_files = cached_file_paths ,
322
+ use_hf_secret = use_hf_secret ,
323
+ hf_access_token_file_name = HF_ACCESS_TOKEN_FILE_NAME ,
315
324
)
316
325
dockerfile_filepath = build_dir / "Dockerfile"
317
326
dockerfile_filepath .write_text (dockerfile_content )
@@ -361,10 +370,10 @@ def prepare_image_build_dir(
361
370
build_dir = build_truss_target_directory (model_framework_name )
362
371
363
372
if config .build .model_server is ModelServer .TGI :
364
- create_tgi_build_dir (config , build_dir , truss_dir )
373
+ create_tgi_build_dir (config , build_dir , truss_dir , use_hf_secret )
365
374
return
366
375
elif config .build .model_server is ModelServer .VLLM :
367
- create_vllm_build_dir (config , build_dir , truss_dir )
376
+ create_vllm_build_dir (config , build_dir , truss_dir , use_hf_secret )
368
377
return
369
378
elif config .build .model_server is ModelServer .TRITON :
370
379
create_triton_build_dir (config , build_dir , truss_dir )
@@ -481,6 +490,7 @@ def _render_dockerfile(
481
490
build_dir / REQUIREMENTS_TXT_FILENAME
482
491
)
483
492
493
+ hf_access_token = config .secrets .get (HF_ACCESS_TOKEN_SECRET_NAME )
484
494
dockerfile_contents = dockerfile_template .render (
485
495
should_install_server_requirements = should_install_server_requirements ,
486
496
base_image_name_and_tag = base_image_name_and_tag ,
@@ -497,6 +507,8 @@ def _render_dockerfile(
497
507
cached_files = cached_files ,
498
508
credentials_exists = credentials_file .exists (),
499
509
hf_cache = len (config .hf_cache .models ) > 0 ,
510
+ hf_access_token = hf_access_token ,
511
+ hf_access_token_file_name = HF_ACCESS_TOKEN_FILE_NAME ,
500
512
)
501
513
docker_file_path = build_dir / MODEL_DOCKERFILE_NAME
502
514
docker_file_path .write_text (dockerfile_contents )
0 commit comments