diff --git a/src/filesystem/api.cc b/src/filesystem/api.cc index 4f3883ccc..2734eaf48 100644 --- a/src/filesystem/api.cc +++ b/src/filesystem/api.cc @@ -619,7 +619,16 @@ MakeTemporaryDirectory(const FileSystemType type, std::string* temp_dir) { std::shared_ptr fs; RETURN_IF_ERROR(fsm_.GetFileSystem(type, fs)); - return fs->MakeTemporaryDirectory(temp_dir); + return fs->MakeTemporaryDirectory(kDefaultMountDirectory, temp_dir); +} + +Status +MakeTemporaryDirectory( + const FileSystemType type, std::string dir_path, std::string* temp_dir) +{ + std::shared_ptr fs; + RETURN_IF_ERROR(fsm_.GetFileSystem(type, fs)); + return fs->MakeTemporaryDirectory(dir_path, temp_dir); } Status diff --git a/src/filesystem/api.h b/src/filesystem/api.h index e7711ce8e..8ee0b373e 100644 --- a/src/filesystem/api.h +++ b/src/filesystem/api.h @@ -198,6 +198,15 @@ Status MakeDirectory(const std::string& dir, const bool recursive); /// \return Error status Status MakeTemporaryDirectory(const FileSystemType type, std::string* temp_dir); +/// Create a temporary directory of the specified filesystem type +/// in a provided location. +/// \param type The type of the filesystem. +/// \param dir_path The specified path. +/// \param temp_dir Returns the path to the temporary directory. +/// \return Error status +Status MakeTemporaryDirectory( + const FileSystemType type, std::string dir_path, std::string* temp_dir); + /// Delete a path. /// \param path The path to the directory or file. /// \return Error status diff --git a/src/filesystem/implementations/as.h b/src/filesystem/implementations/as.h index 13eb80d99..fc449475a 100644 --- a/src/filesystem/implementations/as.h +++ b/src/filesystem/implementations/as.h @@ -94,7 +94,8 @@ class ASFileSystem : public FileSystem { const std::string& path, const char* contents, const size_t content_len) override; Status MakeDirectory(const std::string& dir, const bool recursive) override; - Status MakeTemporaryDirectory(std::string* temp_dir) override; + Status MakeTemporaryDirectory( + std::string dir_path, std::string* temp_dir) override; Status DeletePath(const std::string& path) override; private: @@ -441,17 +442,20 @@ ASFileSystem::LocalizePath( "AS file localization not yet implemented " + path); } - std::string folder_template = "/tmp/folderXXXXXX"; - char* tmp_folder = mkdtemp(const_cast(folder_template.c_str())); - if (tmp_folder == nullptr) { - return Status( - Status::Code::INTERNAL, - "Failed to create local temp folder: " + folder_template + - ", errno:" + strerror(errno)); - } + // Create a local directory for azure model store. + // If ENV variable are not set, creates a temporary directory + // under `/tmp` with the format: "folderXXXXXX". + // Otherwise, will create a folder under specified directory with the same + // format. + std::string env_mount_dir = GetEnvironmentVariableOrDefault( + "TRITON_AZURE_MOUNT_DIRECTORY", kDefaultMountDirectory); + std::string tmp_folder; + RETURN_IF_ERROR(triton::core::MakeTemporaryDirectory( + FileSystemType::LOCAL, env_mount_dir, &tmp_folder)); + localized->reset(new LocalizedPath(path, tmp_folder)); - std::string dest(folder_template); + std::string dest(tmp_folder); std::string container, blob; RETURN_IF_ERROR(ParsePath(path, &container, &blob)); @@ -495,7 +499,8 @@ ASFileSystem::MakeDirectory(const std::string& dir, const bool recursive) } Status -ASFileSystem::MakeTemporaryDirectory(std::string* temp_dir) +ASFileSystem::MakeTemporaryDirectory( + std::string dir_path, std::string* temp_dir) { return Status( Status::Code::UNSUPPORTED, diff --git a/src/filesystem/implementations/common.h b/src/filesystem/implementations/common.h index 5f45444dd..54afb284b 100644 --- a/src/filesystem/implementations/common.h +++ b/src/filesystem/implementations/common.h @@ -65,6 +65,9 @@ namespace triton { namespace core { +// Default folder for temporary local cache +constexpr char kDefaultMountDirectory[] = "/tmp"; + // FileSystem interface that all file system implementation should inherit from. // To add new file system support, the implementation should be added and made // visible to FileSystemManager in api.cc @@ -91,7 +94,8 @@ class FileSystem { const size_t content_len) = 0; virtual Status MakeDirectory( const std::string& dir, const bool recursive) = 0; - virtual Status MakeTemporaryDirectory(std::string* temp_dir) = 0; + virtual Status MakeTemporaryDirectory( + std::string dir_path, std::string* temp_dir) = 0; virtual Status DeletePath(const std::string& path) = 0; }; @@ -106,4 +110,18 @@ AppendSlash(const std::string& name) return (name + "/"); } +/// Helper function to get the value of the environment variable, +/// or default value if not set. +/// +/// \param variable_name The name of the environment variable. +/// \param default_value The default value. +/// \return The environment variable or the default value if not set. +std::string +GetEnvironmentVariableOrDefault( + const std::string& variable_name, const std::string& default_value) +{ + const char* value = getenv(variable_name.c_str()); + return value ? value : default_value; +} + }} // namespace triton::core diff --git a/src/filesystem/implementations/gcs.h b/src/filesystem/implementations/gcs.h index 56c3d8d34..0ffda004d 100644 --- a/src/filesystem/implementations/gcs.h +++ b/src/filesystem/implementations/gcs.h @@ -83,7 +83,8 @@ class GCSFileSystem : public FileSystem { const std::string& path, const char* contents, const size_t content_len) override; Status MakeDirectory(const std::string& dir, const bool recursive) override; - Status MakeTemporaryDirectory(std::string* temp_dir) override; + Status MakeTemporaryDirectory( + std::string dir_path, std::string* temp_dir) override; Status DeletePath(const std::string& path) override; private: @@ -380,9 +381,12 @@ GCSFileSystem::LocalizePath( "GCS file localization not yet implemented " + path); } + // Create a local directory for GCS model store. + std::string env_mount_dir = GetEnvironmentVariableOrDefault( + "TRITON_GCS_MOUNT_DIRECTORY", kDefaultMountDirectory); std::string tmp_folder; - RETURN_IF_ERROR( - triton::core::MakeTemporaryDirectory(FileSystemType::LOCAL, &tmp_folder)); + RETURN_IF_ERROR(triton::core::MakeTemporaryDirectory( + FileSystemType::LOCAL, env_mount_dir, &tmp_folder)); localized->reset(new LocalizedPath(path, tmp_folder)); @@ -480,7 +484,8 @@ GCSFileSystem::MakeDirectory(const std::string& dir, const bool recursive) } Status -GCSFileSystem::MakeTemporaryDirectory(std::string* temp_dir) +GCSFileSystem::MakeTemporaryDirectory( + std::string dir_path, std::string* temp_dir) { return Status( Status::Code::UNSUPPORTED, diff --git a/src/filesystem/implementations/local.h b/src/filesystem/implementations/local.h index f6deeb5e6..8e60054a6 100644 --- a/src/filesystem/implementations/local.h +++ b/src/filesystem/implementations/local.h @@ -57,7 +57,8 @@ class LocalFileSystem : public FileSystem { const std::string& path, const char* contents, const size_t content_len) override; Status MakeDirectory(const std::string& dir, const bool recursive) override; - Status MakeTemporaryDirectory(std::string* temp_dir) override; + Status MakeTemporaryDirectory( + std::string dir_path, std::string* temp_dir) override; Status DeletePath(const std::string& path) override; }; @@ -280,7 +281,8 @@ LocalFileSystem::MakeDirectory(const std::string& dir, const bool recursive) } Status -LocalFileSystem::MakeTemporaryDirectory(std::string* temp_dir) +LocalFileSystem::MakeTemporaryDirectory( + std::string dir_path, std::string* temp_dir) { #ifdef _WIN32 char temp_path[MAX_PATH + 1]; @@ -314,7 +316,10 @@ LocalFileSystem::MakeTemporaryDirectory(std::string* temp_dir) "Failed to create local temp folder: " + *temp_dir); } #else - std::string folder_template = "/tmp/folderXXXXXX"; + if (dir_path.empty()) { + dir_path = kDefaultMountDirectory; + } + std::string folder_template = JoinPath({dir_path, "folderXXXXXX"}); char* res = mkdtemp(const_cast(folder_template.c_str())); if (res == nullptr) { return Status( diff --git a/src/filesystem/implementations/s3.h b/src/filesystem/implementations/s3.h index fe6da18c3..48e1cbc27 100644 --- a/src/filesystem/implementations/s3.h +++ b/src/filesystem/implementations/s3.h @@ -159,7 +159,8 @@ class S3FileSystem : public FileSystem { const std::string& path, const char* contents, const size_t content_len) override; Status MakeDirectory(const std::string& dir, const bool recursive) override; - Status MakeTemporaryDirectory(std::string* temp_dir) override; + Status MakeTemporaryDirectory( + std::string dir_path, std::string* temp_dir) override; Status DeletePath(const std::string& path) override; private: @@ -652,10 +653,16 @@ S3FileSystem::LocalizePath( effective_path = path; } - // Create temporary directory + // Create a local directory for AWS model store. + // If ENV variable are not set, creates a temporary directory + // under `/tmp` with the format: "folderXXXXXX". + // Otherwise, will create a folder under specified directory with the same + // format. + std::string env_mount_dir = GetEnvironmentVariableOrDefault( + "TRITON_AWS_MOUNT_DIRECTORY", kDefaultMountDirectory); std::string tmp_folder; - RETURN_IF_ERROR( - triton::core::MakeTemporaryDirectory(FileSystemType::LOCAL, &tmp_folder)); + RETURN_IF_ERROR(triton::core::MakeTemporaryDirectory( + FileSystemType::LOCAL, env_mount_dir, &tmp_folder)); // Specify contents to be downloaded std::set contents; @@ -774,7 +781,8 @@ S3FileSystem::MakeDirectory(const std::string& dir, const bool recursive) } Status -S3FileSystem::MakeTemporaryDirectory(std::string* temp_dir) +S3FileSystem::MakeTemporaryDirectory( + std::string dir_path, std::string* temp_dir) { return Status( Status::Code::UNSUPPORTED,