Skip to content

Commit 2bd9a25

Browse files
committed
Apply similar changes to gcs
1 parent af94ad7 commit 2bd9a25

File tree

1 file changed

+18
-4
lines changed
  • src/filesystem/implementations

1 file changed

+18
-4
lines changed

src/filesystem/implementations/gcs.h

Lines changed: 18 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -384,9 +384,23 @@ GCSFileSystem::LocalizePath(
384384
"GCS file localization not yet implemented " + path);
385385
}
386386

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");
387393
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+
}
390404

391405
localized->reset(new LocalizedPath(path, tmp_folder));
392406

@@ -406,7 +420,7 @@ GCSFileSystem::LocalizePath(
406420
std::string local_fpath =
407421
JoinPath({(*localized)->Path(), gcs_removed_path});
408422
RETURN_IF_ERROR(IsDirectory(gcs_fpath, &is_subdir));
409-
if (is_subdir) {
423+
if (recursive && is_subdir) {
410424
// Create local mirror of sub-directories
411425
#ifdef _WIN32
412426
int status = mkdir(const_cast<char*>(local_fpath.c_str()));
@@ -429,7 +443,7 @@ GCSFileSystem::LocalizePath(
429443
++itr) {
430444
contents.insert(JoinPath({gcs_fpath, *itr}));
431445
}
432-
} else {
446+
} else if (!is_subdir) {
433447
// Create local copy of file
434448
std::string file_bucket, file_object;
435449
RETURN_IF_ERROR(ParsePath(gcs_fpath, &file_bucket, &file_object));

0 commit comments

Comments
 (0)