Adding ability to specify subdirectory to download for LocalizePath

triton-inference-server · Sep 7, 2023 · e5a40b7 · e5a40b7
1 parent a23786c
commit e5a40b7
Show file tree

Hide file tree

Showing 9 changed files with 53 additions and 35 deletions.
diff --git a/src/backend_model.cc b/src/backend_model.cc
@@ -75,10 +75,11 @@ TritonModel::Create(
   }
 
   // Localize the content of the model repository corresponding to
-  // 'model_path'. This model holds a handle to the localized content
-  // so that it persists as long as the model is loaded.
+  // 'model_path' and model's version. 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, &localized_model_dir));
+  RETURN_IF_ERROR(
+      LocalizePath(model_path, std::to_string(version), &localized_model_dir));
 
   // Localize paths in backend model config
   // [FIXME] Remove once a more permanent solution is implemented (DLIS-4211)

diff --git a/src/filesystem/api.cc b/src/filesystem/api.cc
@@ -558,11 +558,13 @@ ReadTextProto(const std::string& path, google::protobuf::Message* msg)
 }
 
 Status
-LocalizePath(const std::string& path, std::shared_ptr<LocalizedPath>* localized)
+LocalizePath(
+    const std::string& path, const std::string& fetch_subdir,
+    std::shared_ptr<LocalizedPath>* localized)
 {
   std::shared_ptr<FileSystem> fs;
   RETURN_IF_ERROR(fsm_.GetFileSystem(path, fs));
-  return fs->LocalizePath(path, localized);
+  return fs->LocalizePath(path, fetch_subdir, localized);
 }
 
 Status

diff --git a/src/filesystem/api.h b/src/filesystem/api.h
@@ -145,11 +145,15 @@ 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 localized Returns the LocalizedPath object
 /// representing the local copy of the path.
 /// \return Error status
 Status LocalizePath(
-    const std::string& path, std::shared_ptr<LocalizedPath>* localized);
+    const std::string& path, const std::string& fetch_subdir,
+    std::shared_ptr<LocalizedPath>* localized);
 
 /// Write a string to a file.
 /// \param path The path of the file.

diff --git a/src/filesystem/implementations/as.h b/src/filesystem/implementations/as.h
@@ -86,7 +86,7 @@ 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& path, const std::string& fetch_subdir,
       std::shared_ptr<LocalizedPath>* localized) override;
   Status WriteTextFile(
       const std::string& path, const std::string& contents) override;
@@ -424,7 +424,8 @@ ASFileSystem::DownloadFolder(
 
 Status
 ASFileSystem::LocalizePath(
-    const std::string& path, std::shared_ptr<LocalizedPath>* localized)
+    const std::string& path, const std::string& fetch_subdir,
+    std::shared_ptr<LocalizedPath>* localized)
 {
   bool exists;
   RETURN_IF_ERROR(FileExists(path, &exists));

diff --git a/src/filesystem/implementations/common.h b/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, std::shared_ptr<LocalizedPath>* localized) = 0;
+      const std::string& path, const std::string& fetch_subdir,
+      std::shared_ptr<LocalizedPath>* localized) = 0;
   virtual Status WriteTextFile(
       const std::string& path, const std::string& contents) = 0;
   virtual Status WriteBinaryFile(

diff --git a/src/filesystem/implementations/gcs.h b/src/filesystem/implementations/gcs.h
@@ -75,7 +75,7 @@ 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& path, const std::string& fetch_subdir,
       std::shared_ptr<LocalizedPath>* localized) override;
   Status WriteTextFile(
       const std::string& path, const std::string& contents) override;
@@ -363,7 +363,8 @@ GCSFileSystem::ReadTextFile(const std::string& path, std::string* contents)
 
 Status
 GCSFileSystem::LocalizePath(
-    const std::string& path, std::shared_ptr<LocalizedPath>* localized)
+    const std::string& path, const std::string& fetch_subdir,
+    std::shared_ptr<LocalizedPath>* localized)
 {
   bool exists;
   RETURN_IF_ERROR(FileExists(path, &exists));

diff --git a/src/filesystem/implementations/local.h b/src/filesystem/implementations/local.h
@@ -49,7 +49,7 @@ 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& path, const std::string& fetch_subdir,
       std::shared_ptr<LocalizedPath>* localized) override;
   Status WriteTextFile(
       const std::string& path, const std::string& contents) override;
@@ -204,7 +204,8 @@ LocalFileSystem::ReadTextFile(const std::string& path, std::string* contents)
 
 Status
 LocalFileSystem::LocalizePath(
-    const std::string& path, std::shared_ptr<LocalizedPath>* localized)
+    const std::string& path, const std::string& fetch_subdir,
+    std::shared_ptr<LocalizedPath>* localized)
 {
   // For local file system we don't actually need to download the
   // directory or file. We use it in place.

diff --git a/src/filesystem/implementations/s3.h b/src/filesystem/implementations/s3.h
@@ -151,7 +151,7 @@ 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& path, const std::string& fetch_subdir,
       std::shared_ptr<LocalizedPath>* localized) override;
   Status WriteTextFile(
       const std::string& path, const std::string& contents) override;
@@ -628,7 +628,8 @@ S3FileSystem::ReadTextFile(const std::string& path, std::string* contents)
 
 Status
 S3FileSystem::LocalizePath(
-    const std::string& path, std::shared_ptr<LocalizedPath>* localized)
+    const std::string& path, const std::string& fetch_subdir,
+    std::shared_ptr<LocalizedPath>* localized)
 {
   // Check if the directory or file exists
   bool exists;
@@ -693,28 +694,34 @@ 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) {
-        // Create local mirror of sub-directories
+        if (copy_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 {
         // Create local copy of file

diff --git a/src/model_config_utils.cc b/src/model_config_utils.cc
@@ -918,8 +918,8 @@ LocalizePythonBackendExecutionEnvironmentPath(
           model_path_slash) {
         // Localize the file
         std::shared_ptr<LocalizedPath> localized_exec_env_path;
-        RETURN_IF_ERROR(
-            LocalizePath(abs_exec_env_path, &localized_exec_env_path));
+        RETURN_IF_ERROR(LocalizePath(
+            abs_exec_env_path, "" /*fetch_subdir*/, &localized_exec_env_path));
         // Persist the localized temporary path
         (*localized_model_dir)
             ->other_localized_path.push_back(localized_exec_env_path);