Skip to content

Commit af3f71f

Browse files
oandreeva-nvoandreeva-nv
committed
Add ability to specify directory for cloud storage
1 parent 80fef04 commit af3f71f

File tree

7 files changed

+81
-25
lines changed

7 files changed

+81
-25
lines changed

‎src/filesystem/api.cc‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -619,7 +619,16 @@ 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(
627+
const FileSystemType type, std::string& dir_path, std::string* temp_dir)
628+
{
629+
std::shared_ptr<FileSystem> fs;
630+
RETURN_IF_ERROR(fsm_.GetFileSystem(type, fs));
631+
return fs->MakeTemporaryDirectory(dir_path, temp_dir);
623632
}
624633

625634
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: 16 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -94,7 +94,8 @@ 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(
98+
std::string& dir_path, std::string* temp_dir) override;
9899
Status DeletePath(const std::string& path) override;
99100

100101
private:
@@ -441,17 +442,20 @@ ASFileSystem::LocalizePath(
441442
"AS file localization not yet implemented " + path);
442443
}
443444

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));
451-
}
445+
// Create a local directory for azure model store.
446+
// If ENV variable are not set, creates a temporary directory
447+
// under `/tmp` with the format: "folderXXXXXX".
448+
// Otherwise, will create a folder under specified directory with the same
449+
// format.
450+
std::string env_mount_dir =
451+
GetEnvironmentVariableOrDefault("TRITON_AZURE_MOUNT_DIRECTORY", "/tmp");
452+
std::string tmp_folder;
453+
RETURN_IF_ERROR(triton::core::MakeTemporaryDirectory(
454+
FileSystemType::LOCAL, std::string(env_mount_dir), &tmp_folder));
455+
452456
localized->reset(new LocalizedPath(path, tmp_folder));
453457

454-
std::string dest(folder_template);
458+
std::string dest(tmp_folder);
455459

456460
std::string container, blob;
457461
RETURN_IF_ERROR(ParsePath(path, &container, &blob));
@@ -495,7 +499,8 @@ ASFileSystem::MakeDirectory(const std::string& dir, const bool recursive)
495499
}
496500

497501
Status
498-
ASFileSystem::MakeTemporaryDirectory(std::string* temp_dir)
502+
ASFileSystem::MakeTemporaryDirectory(
503+
std::string& dir_path, std::string* temp_dir)
499504
{
500505
return Status(
501506
Status::Code::UNSUPPORTED,

‎src/filesystem/implementations/common.h‎

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -91,7 +91,8 @@ 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(
95+
std::string dir_path, std::string* temp_dir) = 0;
9596
virtual Status DeletePath(const std::string& path) = 0;
9697
};
9798

@@ -106,4 +107,18 @@ AppendSlash(const std::string& name)
106107
return (name + "/");
107108
}
108109

110+
/// Helper function to get the value of the environment variable,
111+
/// or default value if not set.
112+
///
113+
/// \param variable_name The name of the environment variable.
114+
/// \param default_value The default value.
115+
/// \return The environment variable or the default value if not set.
116+
std::string
117+
GetEnvironmentVariableOrDefault(
118+
const std::string& variable_name, const std::string& default_value)
119+
{
120+
const char* value = getenv(variable_name.c_str());
121+
return value ? value : default_value;
122+
}
123+
109124
}} // namespace triton::core

‎src/filesystem/implementations/gcs.h‎

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -83,7 +83,8 @@ 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(
87+
std::string& dir_path, std::string* temp_dir) override;
8788
Status DeletePath(const std::string& path) override;
8889

8990
private:
@@ -380,9 +381,12 @@ GCSFileSystem::LocalizePath(
380381
"GCS file localization not yet implemented " + path);
381382
}
382383

384+
// Create a local directory for GCS model store.
385+
std::string env_mount_dir =
386+
GetEnvironmentVariableOrDefault("TRITON_GCS_MOUNT_DIRECTORY", "/tmp");
383387
std::string tmp_folder;
384-
RETURN_IF_ERROR(
385-
triton::core::MakeTemporaryDirectory(FileSystemType::LOCAL, &tmp_folder));
388+
RETURN_IF_ERROR(triton::core::MakeTemporaryDirectory(
389+
FileSystemType::LOCAL, std::string(env_mount_dir), &tmp_folder));
386390

387391
localized->reset(new LocalizedPath(path, tmp_folder));
388392

@@ -480,7 +484,8 @@ GCSFileSystem::MakeDirectory(const std::string& dir, const bool recursive)
480484
}
481485

482486
Status
483-
GCSFileSystem::MakeTemporaryDirectory(std::string* temp_dir)
487+
GCSFileSystem::MakeTemporaryDirectory(
488+
std::string& dir_path, std::string* temp_dir)
484489
{
485490
return Status(
486491
Status::Code::UNSUPPORTED,

‎src/filesystem/implementations/local.h‎

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -57,7 +57,8 @@ 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(
61+
std::string& dir_path, std::string* temp_dir) override;
6162
Status DeletePath(const std::string& path) override;
6263
};
6364

@@ -280,7 +281,8 @@ LocalFileSystem::MakeDirectory(const std::string& dir, const bool recursive)
280281
}
281282

282283
Status
283-
LocalFileSystem::MakeTemporaryDirectory(std::string* temp_dir)
284+
LocalFileSystem::MakeTemporaryDirectory(
285+
std::string& dir_path, std::string* temp_dir)
284286
{
285287
#ifdef _WIN32
286288
char temp_path[MAX_PATH + 1];
@@ -314,7 +316,10 @@ LocalFileSystem::MakeTemporaryDirectory(std::string* temp_dir)
314316
"Failed to create local temp folder: " + *temp_dir);
315317
}
316318
#else
317-
std::string folder_template = "/tmp/folderXXXXXX";
319+
if (dir_path.empty()) {
320+
dir_path = "/tmp";
321+
}
322+
std::string folder_template = JoinPath({dir_path, "folderXXXXXX"});
318323
char* res = mkdtemp(const_cast<char*>(folder_template.c_str()));
319324
if (res == nullptr) {
320325
return Status(

‎src/filesystem/implementations/s3.h‎

Lines changed: 13 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -159,7 +159,8 @@ 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(
163+
std::string& dir_path, std::string* temp_dir) override;
163164
Status DeletePath(const std::string& path) override;
164165

165166
private:
@@ -652,10 +653,16 @@ S3FileSystem::LocalizePath(
652653
effective_path = path;
653654
}
654655

655-
// Create temporary directory
656+
// Create a local directory for AWS model store.
657+
// If ENV variable are not set, creates a temporary directory
658+
// under `/tmp` with the format: "folderXXXXXX".
659+
// Otherwise, will create a folder under specified directory with the same
660+
// format.
661+
std::string env_mount_dir =
662+
GetEnvironmentVariableOrDefault("TRITON_AWS_MOUNT_DIRECTORY", "/tmp");
656663
std::string tmp_folder;
657-
RETURN_IF_ERROR(
658-
triton::core::MakeTemporaryDirectory(FileSystemType::LOCAL, &tmp_folder));
664+
RETURN_IF_ERROR(triton::core::MakeTemporaryDirectory(
665+
FileSystemType::LOCAL, std::string(env_mount_dir), &tmp_folder));
659666

660667
// Specify contents to be downloaded
661668
std::set<std::string> contents;
@@ -774,7 +781,8 @@ S3FileSystem::MakeDirectory(const std::string& dir, const bool recursive)
774781
}
775782

776783
Status
777-
S3FileSystem::MakeTemporaryDirectory(std::string* temp_dir)
784+
S3FileSystem::MakeTemporaryDirectory(
785+
std::string& dir_path, std::string* temp_dir)
778786
{
779787
return Status(
780788
Status::Code::UNSUPPORTED,

0 commit comments

Comments
 (0)