Generalized implementation of : added recursive flag and mount_dir

Author: oandreeva-nv

src/backend_model.cc

@@ -75,11 +75,19 @@ TritonModel::Create(
   }
 
   // Localize the content of the model repository corresponding to
-  // 'model_path' and model's version. This model holds a handle to
+  // 'model_path'. This model holds a handle to
   // the localized content so that it persists as long as the model is loaded.
   std::shared_ptr<LocalizedPath> localized_model_dir;
   RETURN_IF_ERROR(
-      LocalizePath(model_path, std::to_string(version), &localized_model_dir));
+      LocalizePath(model_path, false /*recursive*/, &localized_model_dir));
+
+  const auto version_path = JoinPath({model_path, std::to_string(version)});
+  std::shared_ptr<LocalizedPath> localized_version_dir;
+  RETURN_IF_ERROR(LocalizePath(
+      version_path, true /*recursive*/, localized_model_dir->Path(),
+      &localized_version_dir));
+
+  localized_model_dir->other_localized_path.push_back(localized_version_dir);
 
   // Localize paths in backend model config
   // [FIXME] Remove once a more permanent solution is implemented (DLIS-4211)
@@ -111,12 +119,11 @@ TritonModel::Create(
   // Get the path to the backend shared library. Search path is
   // version directory, model directory, global backend directory.
   const auto localized_model_path = localized_model_dir->Path();
-  const auto version_path =
-      JoinPath({localized_model_path, std::to_string(version)});
+  const auto localized_version_path = localized_version_dir->Path();
   const std::string global_path =
       JoinPath({backend_dir, specialized_backend_name});
   const std::vector<std::string> search_paths = {
-      version_path, localized_model_path, global_path};
+      localized_version_path, localized_model_path, global_path};
 
   std::string backend_libdir;
   std::string backend_libpath;

src/filesystem/api.cc

@@ -559,12 +559,22 @@ ReadTextProto(const std::string& path, google::protobuf::Message* msg)
 
 Status
 LocalizePath(
-    const std::string& path, const std::string& fetch_subdir,
+    const std::string& path, const bool recursive,
     std::shared_ptr<LocalizedPath>* localized)
 {
   std::shared_ptr<FileSystem> fs;
   RETURN_IF_ERROR(fsm_.GetFileSystem(path, fs));
-  return fs->LocalizePath(path, fetch_subdir, localized);
+  return fs->LocalizePath(path, recursive, "", localized);
+}
+
+Status
+LocalizePath(
+    const std::string& path, const bool recursive, const std::string& mount_dir,
+    std::shared_ptr<LocalizedPath>* localized)
+{
+  std::shared_ptr<FileSystem> fs;
+  RETURN_IF_ERROR(fsm_.GetFileSystem(path, fs));
+  return fs->LocalizePath(path, recursive, mount_dir, localized);
 }
 
 Status

src/filesystem/api.h

@@ -145,14 +145,26 @@ Status ReadTextFile(const std::string& path, std::string* contents);
 
 /// Create an object representing a local copy of a path.
 /// \param path The path of the directory or file.
-/// \param fetch_subdir If specified, will only download provided
-/// sub directory, otherwise all subdirectories will be downloaded.
-/// Does not affect files individual files, located under `path`.
+/// \param recursive If true, will fetch all sub-directories in
+/// the provided path.
 /// \param localized Returns the LocalizedPath object
 /// representing the local copy of the path.
 /// \return Error status
 Status LocalizePath(
-    const std::string& path, const std::string& fetch_subdir,
+    const std::string& path, const bool recursive,
+    std::shared_ptr<LocalizedPath>* localized);
+
+/// Create an object representing a local copy of a path.
+/// \param path The path of the directory or file.
+/// \param recursive If true, will fetch all sub-directories in
+/// the provided path.
+/// \param mount_dir If specified, will use provided local directory
+/// for localization.
+/// \param localized Returns the LocalizedPath object
+/// representing the local copy of the path.
+/// \return Error status
+Status LocalizePath(
+    const std::string& path, const bool recursive, const std::string& mount_dir,
     std::shared_ptr<LocalizedPath>* localized);
 
 /// Write a string to a file.

src/filesystem/implementations/as.h

@@ -86,7 +86,8 @@ class ASFileSystem : public FileSystem {
       const std::string& path, std::set<std::string>* files) override;
   Status ReadTextFile(const std::string& path, std::string* contents) override;
   Status LocalizePath(
-      const std::string& path, const std::string& fetch_subdir,
+      const std::string& path, const bool recursive,
+      const std::string& mount_dir,
       std::shared_ptr<LocalizedPath>* localized) override;
   Status WriteTextFile(
       const std::string& path, const std::string& contents) override;
@@ -424,7 +425,7 @@ ASFileSystem::DownloadFolder(
 
 Status
 ASFileSystem::LocalizePath(
-    const std::string& path, const std::string& fetch_subdir,
+    const std::string& path, const bool recursive, const std::string& mount_dir,
     std::shared_ptr<LocalizedPath>* localized)
 {
   bool exists;

src/filesystem/implementations/common.h

@@ -83,7 +83,8 @@ class FileSystem {
   virtual Status ReadTextFile(
       const std::string& path, std::string* contents) = 0;
   virtual Status LocalizePath(
-      const std::string& path, const std::string& fetch_subdir,
+      const std::string& path, const bool recursive,
+      const std::string& mount_dir,
       std::shared_ptr<LocalizedPath>* localized) = 0;
   virtual Status WriteTextFile(
       const std::string& path, const std::string& contents) = 0;

src/filesystem/implementations/gcs.h

@@ -75,7 +75,8 @@ class GCSFileSystem : public FileSystem {
       const std::string& path, std::set<std::string>* files) override;
   Status ReadTextFile(const std::string& path, std::string* contents) override;
   Status LocalizePath(
-      const std::string& path, const std::string& fetch_subdir,
+      const std::string& path, const bool recursive,
+      const std::string& mount_dir,
       std::shared_ptr<LocalizedPath>* localized) override;
   Status WriteTextFile(
       const std::string& path, const std::string& contents) override;
@@ -363,7 +364,7 @@ GCSFileSystem::ReadTextFile(const std::string& path, std::string* contents)
 
 Status
 GCSFileSystem::LocalizePath(
-    const std::string& path, const std::string& fetch_subdir,
+    const std::string& path, const bool recursive, const std::string& mount_dir,
     std::shared_ptr<LocalizedPath>* localized)
 {
   bool exists;

src/filesystem/implementations/local.h

@@ -49,7 +49,8 @@ class LocalFileSystem : public FileSystem {
       const std::string& path, std::set<std::string>* files) override;
   Status ReadTextFile(const std::string& path, std::string* contents) override;
   Status LocalizePath(
-      const std::string& path, const std::string& fetch_subdir,
+      const std::string& path, const bool recursive,
+      const std::string& mount_dir,
       std::shared_ptr<LocalizedPath>* localized) override;
   Status WriteTextFile(
       const std::string& path, const std::string& contents) override;
@@ -204,7 +205,7 @@ LocalFileSystem::ReadTextFile(const std::string& path, std::string* contents)
 
 Status
 LocalFileSystem::LocalizePath(
-    const std::string& path, const std::string& fetch_subdir,
+    const std::string& path, const bool recursive, const std::string& mount_dir,
     std::shared_ptr<LocalizedPath>* localized)
 {
   // For local file system we don't actually need to download the
@@ -255,8 +256,12 @@ LocalFileSystem::MakeDirectory(const std::string& dir, const bool recursive)
   if (mkdir(dir.c_str(), S_IRWXU) == -1)
 #endif
   {
-    // Only allow the error due to parent directory does not exist
-    // if 'recursive' is requested
+    // Return success if directory already exists
+    if (errno == EEXIST) {
+      return Status::Success;
+    }
+    // In all other cases only allow the error due to parent directory
+    // does not exist, if 'recursive' is requested
     if ((errno == ENOENT) && (!dir.empty()) && recursive) {
       RETURN_IF_ERROR(MakeDirectory(DirName(dir), recursive));
       // Retry the creation

src/filesystem/implementations/s3.h

@@ -151,7 +151,8 @@ class S3FileSystem : public FileSystem {
       const std::string& path, std::set<std::string>* files) override;
   Status ReadTextFile(const std::string& path, std::string* contents) override;
   Status LocalizePath(
-      const std::string& path, const std::string& fetch_subdir,
+      const std::string& path, const bool recursive,
+      const std::string& mount_dir,
       std::shared_ptr<LocalizedPath>* localized) override;
   Status WriteTextFile(
       const std::string& path, const std::string& contents) override;
@@ -628,7 +629,7 @@ S3FileSystem::ReadTextFile(const std::string& path, std::string* contents)
 
 Status
 S3FileSystem::LocalizePath(
-    const std::string& path, const std::string& fetch_subdir,
+    const std::string& path, const bool recursive, const std::string& mount_dir,
     std::shared_ptr<LocalizedPath>* localized)
 {
   // Check if the directory or file exists
@@ -653,10 +654,18 @@ S3FileSystem::LocalizePath(
     effective_path = path;
   }
 
-  // Create temporary directory
+  // Create a local directory for s3 model store
+  const char* env_mount_dir = std::getenv("TRITON_AWS_MOUNT_DIRECTORY");
   std::string tmp_folder;
-  RETURN_IF_ERROR(
-      triton::core::MakeTemporaryDirectory(FileSystemType::LOCAL, &tmp_folder));
+  if (mount_dir.empty() && env_mount_dir == nullptr) {
+    RETURN_IF_ERROR(triton::core::MakeTemporaryDirectory(
+        FileSystemType::LOCAL, &tmp_folder));
+  } else {
+    tmp_folder = mount_dir.empty() ? std::string(env_mount_dir) : mount_dir;
+    tmp_folder =
+        JoinPath({tmp_folder, path.substr(path.find_last_of('/') + 1)});
+    RETURN_IF_ERROR(triton::core::MakeDirectory(tmp_folder, true));
+  }
 
   // Specify contents to be downloaded
   std::set<std::string> contents;
@@ -694,36 +703,30 @@ S3FileSystem::LocalizePath(
               : JoinPath({(*localized)->Path(), s3_removed_path});
       bool is_subdir;
       RETURN_IF_ERROR(IsDirectory(s3_fpath, &is_subdir));
-      bool copy_subdir =
-          !fetch_subdir.empty()
-              ? s3_fpath == JoinPath({effective_path, fetch_subdir})
-              : true;
-      if (is_subdir) {
-        if (copy_subdir) {
-          // Create local mirror of sub-directories
+      if (recursive && is_subdir) {
+        // Create local mirror of sub-directories
 #ifdef _WIN32
-          int status = mkdir(const_cast<char*>(local_fpath.c_str()));
+        int status = mkdir(const_cast<char*>(local_fpath.c_str()));
 #else
-          int status = mkdir(
-              const_cast<char*>(local_fpath.c_str()),
-              S_IRUSR | S_IWUSR | S_IXUSR);
+        int status = mkdir(
+            const_cast<char*>(local_fpath.c_str()),
+            S_IRUSR | S_IWUSR | S_IXUSR);
 #endif
-          if (status == -1) {
-            return Status(
-                Status::Code::INTERNAL,
-                "Failed to create local folder: " + local_fpath +
-                    ", errno:" + strerror(errno));
-          }
-
-          // Add sub-directories and deeper files to contents
-          std::set<std::string> subdir_contents;
-          RETURN_IF_ERROR(GetDirectoryContents(s3_fpath, &subdir_contents));
-          for (auto itr = subdir_contents.begin(); itr != subdir_contents.end();
-               ++itr) {
-            contents.insert(JoinPath({s3_fpath, *itr}));
-          }
+        if (status == -1) {
+          return Status(
+              Status::Code::INTERNAL,
+              "Failed to create local folder: " + local_fpath +
+                  ", errno:" + strerror(errno));
+        }
+
+        // Add sub-directories and deeper files to contents
+        std::set<std::string> subdir_contents;
+        RETURN_IF_ERROR(GetDirectoryContents(s3_fpath, &subdir_contents));
+        for (auto itr = subdir_contents.begin(); itr != subdir_contents.end();
+             ++itr) {
+          contents.insert(JoinPath({s3_fpath, *itr}));
         }
-      } else {
+      } else if (!is_subdir) {
         // Create local copy of file
         std::string file_bucket, file_object;
         RETURN_IF_ERROR(ParsePath(s3_fpath, &file_bucket, &file_object));

src/model_config_utils.cc

@@ -919,7 +919,7 @@ LocalizePythonBackendExecutionEnvironmentPath(
         // Localize the file
         std::shared_ptr<LocalizedPath> localized_exec_env_path;
         RETURN_IF_ERROR(LocalizePath(
-            abs_exec_env_path, "" /*fetch_subdir*/, &localized_exec_env_path));
+            abs_exec_env_path, true /*recursive*/, &localized_exec_env_path));
         // Persist the localized temporary path
         (*localized_model_dir)
             ->other_localized_path.push_back(localized_exec_env_path);