Specify non-random location for filesystem's temporary cache

src/filesystem/api.cc

@@ -619,7 +619,16 @@ MakeTemporaryDirectory(const FileSystemType type, std::string* temp_dir)
 {
   std::shared_ptr<FileSystem> 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<FileSystem> fs;
+  RETURN_IF_ERROR(fsm_.GetFileSystem(type, fs));
+  return fs->MakeTemporaryDirectory(dir_path, temp_dir);
 }
 
 Status

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

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<char*>(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,

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

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,

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<char*>(folder_template.c_str()));
   if (res == nullptr) {
     return Status(

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<std::string> 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,