@@ -384,9 +384,23 @@ GCSFileSystem::LocalizePath(
384
384
" GCS file localization not yet implemented " + path);
385
385
}
386
386
387
+ // Create a local directory for s3 model store.
388
+ // If `mount_dir` or ENV variable are not set,
389
+ // creates a temporary directory under `/tmp` with the format: "folderXXXXXX".
390
+ // Otherwise, will create a folder under specified directory with the name
391
+ // indicated in path (i.e. everything after the last encounter of `/`).
392
+ const char * env_mount_dir = std::getenv (" TRITON_GCS_MOUNT_DIRECTORY" );
387
393
std::string tmp_folder;
388
- RETURN_IF_ERROR (
389
- triton::core::MakeTemporaryDirectory (FileSystemType::LOCAL, &tmp_folder));
394
+ if (mount_dir.empty () && env_mount_dir == nullptr ) {
395
+ RETURN_IF_ERROR (triton::core::MakeTemporaryDirectory (
396
+ FileSystemType::LOCAL, &tmp_folder));
397
+ } else {
398
+ tmp_folder = mount_dir.empty () ? std::string (env_mount_dir) : mount_dir;
399
+ tmp_folder =
400
+ JoinPath ({tmp_folder, path.substr (path.find_last_of (' /' ) + 1 )});
401
+ RETURN_IF_ERROR (triton::core::MakeDirectory (
402
+ tmp_folder, true /* recursive*/ , true /* allow_dir_exist*/ ));
403
+ }
390
404
391
405
localized->reset (new LocalizedPath (path, tmp_folder));
392
406
@@ -406,7 +420,7 @@ GCSFileSystem::LocalizePath(
406
420
std::string local_fpath =
407
421
JoinPath ({(*localized)->Path (), gcs_removed_path});
408
422
RETURN_IF_ERROR (IsDirectory (gcs_fpath, &is_subdir));
409
- if (is_subdir) {
423
+ if (recursive && is_subdir) {
410
424
// Create local mirror of sub-directories
411
425
#ifdef _WIN32
412
426
int status = mkdir (const_cast <char *>(local_fpath.c_str ()));
@@ -429,7 +443,7 @@ GCSFileSystem::LocalizePath(
429
443
++itr) {
430
444
contents.insert (JoinPath ({gcs_fpath, *itr}));
431
445
}
432
- } else {
446
+ } else if (!is_subdir) {
433
447
// Create local copy of file
434
448
std::string file_bucket, file_object;
435
449
RETURN_IF_ERROR (ParsePath (gcs_fpath, &file_bucket, &file_object));
0 commit comments