Skip to content

Commit 2bbccf5

Browse files
author
oandreeva-nv
committed
Add ability to specify directory for cloud storage
1 parent 80fef04 commit 2bbccf5

File tree

7 files changed

+65
-24
lines changed

7 files changed

+65
-24
lines changed

‎src/filesystem/api.cc‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -619,7 +619,15 @@ MakeTemporaryDirectory(const FileSystemType type, std::string* temp_dir)
619619
{
620620
std::shared_ptr<FileSystem> fs;
621621
RETURN_IF_ERROR(fsm_.GetFileSystem(type, fs));
622-
return fs->MakeTemporaryDirectory(temp_dir);
622+
return fs->MakeTemporaryDirectory("/tmp", temp_dir);
623+
}
624+
625+
Status
626+
MakeTemporaryDirectory(const FileSystemType type, std::string dir_path, std::string* temp_dir)
627+
{
628+
std::shared_ptr<FileSystem> fs;
629+
RETURN_IF_ERROR(fsm_.GetFileSystem(type, fs));
630+
return fs->MakeTemporaryDirectory(dir_path, temp_dir);
623631
}
624632

625633
Status

‎src/filesystem/api.h‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -198,6 +198,15 @@ Status MakeDirectory(const std::string& dir, const bool recursive);
198198
/// \return Error status
199199
Status MakeTemporaryDirectory(const FileSystemType type, std::string* temp_dir);
200200

201+
/// Create a temporary directory of the specified filesystem type
202+
/// in a provided location.
203+
/// \param type The type of the filesystem.
204+
/// \param dir_path The specified path.
205+
/// \param temp_dir Returns the path to the temporary directory.
206+
/// \return Error status
207+
Status MakeTemporaryDirectory(
208+
const FileSystemType type, std::string dir_path, std::string* temp_dir);
209+
201210
/// Delete a path.
202211
/// \param path The path to the directory or file.
203212
/// \return Error status

‎src/filesystem/implementations/as.h‎

Lines changed: 17 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -94,7 +94,7 @@ class ASFileSystem : public FileSystem {
9494
const std::string& path, const char* contents,
9595
const size_t content_len) override;
9696
Status MakeDirectory(const std::string& dir, const bool recursive) override;
97-
Status MakeTemporaryDirectory(std::string* temp_dir) override;
97+
Status MakeTemporaryDirectory(std::string dir_path, std::string* temp_dir) override;
9898
Status DeletePath(const std::string& path) override;
9999

100100
private:
@@ -441,17 +441,24 @@ ASFileSystem::LocalizePath(
441441
"AS file localization not yet implemented " + path);
442442
}
443443

444-
std::string folder_template = "/tmp/folderXXXXXX";
445-
char* tmp_folder = mkdtemp(const_cast<char*>(folder_template.c_str()));
446-
if (tmp_folder == nullptr) {
447-
return Status(
448-
Status::Code::INTERNAL,
449-
"Failed to create local temp folder: " + folder_template +
450-
", errno:" + strerror(errno));
444+
// Create a local directory for azure model store.
445+
// If ENV variable are not set, creates a temporary directory
446+
// under `/tmp` with the format: "folderXXXXXX".
447+
// Otherwise, will create a folder under specified directory with the same
448+
// format.
449+
const char* env_mount_dir = std::getenv("TRITON_AZURE_MOUNT_DIRECTORY");
450+
std::string tmp_folder;
451+
if (env_mount_dir == nullptr) {
452+
RETURN_IF_ERROR(triton::core::MakeTemporaryDirectory(
453+
FileSystemType::LOCAL, &tmp_folder));
454+
} else {
455+
RETURN_IF_ERROR(triton::core::MakeTemporaryDirectory(
456+
FileSystemType::LOCAL, std::string(env_mount_dir), &tmp_folder));
451457
}
458+
452459
localized->reset(new LocalizedPath(path, tmp_folder));
453460

454-
std::string dest(folder_template);
461+
std::string dest(tmp_folder);
455462

456463
std::string container, blob;
457464
RETURN_IF_ERROR(ParsePath(path, &container, &blob));
@@ -495,7 +502,7 @@ ASFileSystem::MakeDirectory(const std::string& dir, const bool recursive)
495502
}
496503

497504
Status
498-
ASFileSystem::MakeTemporaryDirectory(std::string* temp_dir)
505+
ASFileSystem::MakeTemporaryDirectory(std::string dir_path, std::string* temp_dir)
499506
{
500507
return Status(
501508
Status::Code::UNSUPPORTED,

‎src/filesystem/implementations/common.h‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -91,7 +91,7 @@ class FileSystem {
9191
const size_t content_len) = 0;
9292
virtual Status MakeDirectory(
9393
const std::string& dir, const bool recursive) = 0;
94-
virtual Status MakeTemporaryDirectory(std::string* temp_dir) = 0;
94+
virtual Status MakeTemporaryDirectory(std::string dir_path, std::string* temp_dir) = 0;
9595
virtual Status DeletePath(const std::string& path) = 0;
9696
};
9797

‎src/filesystem/implementations/gcs.h‎

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -83,7 +83,7 @@ class GCSFileSystem : public FileSystem {
8383
const std::string& path, const char* contents,
8484
const size_t content_len) override;
8585
Status MakeDirectory(const std::string& dir, const bool recursive) override;
86-
Status MakeTemporaryDirectory(std::string* temp_dir) override;
86+
Status MakeTemporaryDirectory(std::string dir_path, std::string* temp_dir) override;
8787
Status DeletePath(const std::string& path) override;
8888

8989
private:
@@ -380,9 +380,16 @@ GCSFileSystem::LocalizePath(
380380
"GCS file localization not yet implemented " + path);
381381
}
382382

383+
// Create a local directory for GCS model store.
384+
const char* env_mount_dir = std::getenv("TRITON_GCS_MOUNT_DIRECTORY");
383385
std::string tmp_folder;
384-
RETURN_IF_ERROR(
385-
triton::core::MakeTemporaryDirectory(FileSystemType::LOCAL, &tmp_folder));
386+
if (env_mount_dir == nullptr) {
387+
RETURN_IF_ERROR(
388+
triton::core::MakeTemporaryDirectory(FileSystemType::LOCAL, &tmp_folder));
389+
} else {
390+
RETURN_IF_ERROR(
391+
triton::core::MakeTemporaryDirectory(FileSystemType::LOCAL, std::string(env_mount_dir), &tmp_folder));
392+
}
386393

387394
localized->reset(new LocalizedPath(path, tmp_folder));
388395

@@ -480,7 +487,7 @@ GCSFileSystem::MakeDirectory(const std::string& dir, const bool recursive)
480487
}
481488

482489
Status
483-
GCSFileSystem::MakeTemporaryDirectory(std::string* temp_dir)
490+
GCSFileSystem::MakeTemporaryDirectory(std::string dir_path, std::string* temp_dir)
484491
{
485492
return Status(
486493
Status::Code::UNSUPPORTED,

‎src/filesystem/implementations/local.h‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -57,7 +57,7 @@ class LocalFileSystem : public FileSystem {
5757
const std::string& path, const char* contents,
5858
const size_t content_len) override;
5959
Status MakeDirectory(const std::string& dir, const bool recursive) override;
60-
Status MakeTemporaryDirectory(std::string* temp_dir) override;
60+
Status MakeTemporaryDirectory(std::string dir_path, std::string* temp_dir) override;
6161
Status DeletePath(const std::string& path) override;
6262
};
6363

@@ -280,7 +280,7 @@ LocalFileSystem::MakeDirectory(const std::string& dir, const bool recursive)
280280
}
281281

282282
Status
283-
LocalFileSystem::MakeTemporaryDirectory(std::string* temp_dir)
283+
LocalFileSystem::MakeTemporaryDirectory(std::string dir_path, std::string* temp_dir)
284284
{
285285
#ifdef _WIN32
286286
char temp_path[MAX_PATH + 1];
@@ -314,7 +314,7 @@ LocalFileSystem::MakeTemporaryDirectory(std::string* temp_dir)
314314
"Failed to create local temp folder: " + *temp_dir);
315315
}
316316
#else
317-
std::string folder_template = "/tmp/folderXXXXXX";
317+
std::string folder_template = JoinPath({dir_path, "folderXXXXXX"});
318318
char* res = mkdtemp(const_cast<char*>(folder_template.c_str()));
319319
if (res == nullptr) {
320320
return Status(

‎src/filesystem/implementations/s3.h‎

Lines changed: 15 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -159,7 +159,7 @@ class S3FileSystem : public FileSystem {
159159
const std::string& path, const char* contents,
160160
const size_t content_len) override;
161161
Status MakeDirectory(const std::string& dir, const bool recursive) override;
162-
Status MakeTemporaryDirectory(std::string* temp_dir) override;
162+
Status MakeTemporaryDirectory(std::string dir_path, std::string* temp_dir) override;
163163
Status DeletePath(const std::string& path) override;
164164

165165
private:
@@ -652,10 +652,20 @@ S3FileSystem::LocalizePath(
652652
effective_path = path;
653653
}
654654

655-
// Create temporary directory
655+
// Create a local directory for AWS model store.
656+
// If ENV variable are not set, creates a temporary directory
657+
// under `/tmp` with the format: "folderXXXXXX".
658+
// Otherwise, will create a folder under specified directory with the same
659+
// format.
660+
const char* env_mount_dir = std::getenv("TRITON_AWS_MOUNT_DIRECTORY");
656661
std::string tmp_folder;
657-
RETURN_IF_ERROR(
658-
triton::core::MakeTemporaryDirectory(FileSystemType::LOCAL, &tmp_folder));
662+
if (env_mount_dir == nullptr) {
663+
RETURN_IF_ERROR(
664+
triton::core::MakeTemporaryDirectory(FileSystemType::LOCAL, &tmp_folder));
665+
} else {
666+
RETURN_IF_ERROR(
667+
triton::core::MakeTemporaryDirectory(FileSystemType::LOCAL, std::string(env_mount_dir), &tmp_folder));
668+
}
659669

660670
// Specify contents to be downloaded
661671
std::set<std::string> contents;
@@ -774,7 +784,7 @@ S3FileSystem::MakeDirectory(const std::string& dir, const bool recursive)
774784
}
775785

776786
Status
777-
S3FileSystem::MakeTemporaryDirectory(std::string* temp_dir)
787+
S3FileSystem::MakeTemporaryDirectory(std::string dir_path, std::string* temp_dir)
778788
{
779789
return Status(
780790
Status::Code::UNSUPPORTED,

0 commit comments

Comments
 (0)