Skip to content

Commit 23498bc

Browse files
oandreeva-nvoandreeva-nv
committed
Added helper function for environment variable
1 parent 2bbccf5 commit 23498bc

File tree

4 files changed

+41
-32
lines changed

4 files changed

+41
-32
lines changed

‎src/filesystem/implementations/as.h

Lines changed: 8 additions & 10 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 dir_path, 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:
@@ -442,19 +443,15 @@ ASFileSystem::LocalizePath(
442443
}
443444

444445
// Create a local directory for azure model store.
445-
// If ENV variable are not set, creates a temporary directory
446+
// If ENV variable are not set, creates a temporary directory
446447
// under `/tmp` with the format: "folderXXXXXX".
447448
// Otherwise, will create a folder under specified directory with the same
448449
// format.
449-
const char* env_mount_dir = std::getenv("TRITON_AZURE_MOUNT_DIRECTORY");
450+
std::string env_mount_dir =
451+
GetEnvironmentVariableOrDefault("TRITON_AZURE_MOUNT_DIRECTORY", "/tmp");
450452
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(
453+
RETURN_IF_ERROR(triton::core::MakeTemporaryDirectory(
456454
FileSystemType::LOCAL, std::string(env_mount_dir), &tmp_folder));
457-
}
458455

459456
localized->reset(new LocalizedPath(path, tmp_folder));
460457

@@ -502,7 +499,8 @@ ASFileSystem::MakeDirectory(const std::string& dir, const bool recursive)
502499
}
503500

504501
Status
505-
ASFileSystem::MakeTemporaryDirectory(std::string dir_path, std::string* temp_dir)
502+
ASFileSystem::MakeTemporaryDirectory(
503+
std::string dir_path, std::string* temp_dir)
506504
{
507505
return Status(
508506
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 dir_path, 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: 8 additions & 10 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 dir_path, 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:
@@ -381,15 +382,11 @@ GCSFileSystem::LocalizePath(
381382
}
382383

383384
// Create a local directory for GCS model store.
384-
const char* env_mount_dir = std::getenv("TRITON_GCS_MOUNT_DIRECTORY");
385+
std::string env_mount_dir =
386+
GetEnvironmentVariableOrDefault("TRITON_GCS_MOUNT_DIRECTORY", "/tmp");
385387
std::string 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-
}
388+
RETURN_IF_ERROR(triton::core::MakeTemporaryDirectory(
389+
FileSystemType::LOCAL, std::string(env_mount_dir), &tmp_folder));
393390

394391
localized->reset(new LocalizedPath(path, tmp_folder));
395392

@@ -487,7 +484,8 @@ GCSFileSystem::MakeDirectory(const std::string& dir, const bool recursive)
487484
}
488485

489486
Status
490-
GCSFileSystem::MakeTemporaryDirectory(std::string dir_path, std::string* temp_dir)
487+
GCSFileSystem::MakeTemporaryDirectory(
488+
std::string dir_path, std::string* temp_dir)
491489
{
492490
return Status(
493491
Status::Code::UNSUPPORTED,

‎src/filesystem/implementations/s3.h

Lines changed: 9 additions & 11 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 dir_path, 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:
@@ -653,19 +654,15 @@ S3FileSystem::LocalizePath(
653654
}
654655

655656
// Create a local directory for AWS model store.
656-
// If ENV variable are not set, creates a temporary directory
657+
// If ENV variable are not set, creates a temporary directory
657658
// under `/tmp` with the format: "folderXXXXXX".
658659
// Otherwise, will create a folder under specified directory with the same
659660
// format.
660-
const char* env_mount_dir = std::getenv("TRITON_AWS_MOUNT_DIRECTORY");
661+
std::string env_mount_dir =
662+
GetEnvironmentVariableOrDefault("TRITON_AWS_MOUNT_DIRECTORY", "/tmp");
661663
std::string 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-
}
664+
RETURN_IF_ERROR(triton::core::MakeTemporaryDirectory(
665+
FileSystemType::LOCAL, std::string(env_mount_dir), &tmp_folder));
669666

670667
// Specify contents to be downloaded
671668
std::set<std::string> contents;
@@ -784,7 +781,8 @@ S3FileSystem::MakeDirectory(const std::string& dir, const bool recursive)
784781
}
785782

786783
Status
787-
S3FileSystem::MakeTemporaryDirectory(std::string dir_path, std::string* temp_dir)
784+
S3FileSystem::MakeTemporaryDirectory(
785+
std::string dir_path, std::string* temp_dir)
788786
{
789787
return Status(
790788
Status::Code::UNSUPPORTED,

0 commit comments

Comments
 (0)