Skip to content

Commit e5a40b7

Browse files
oandreeva-nvoandreeva-nv
committed
Adding ability to specify subdirectory to download for LocalizePath 
1 parent a23786c commit e5a40b7

File tree

9 files changed

+53
-35
lines changed

9 files changed

+53
-35
lines changed

‎src/backend_model.cc

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -75,10 +75,11 @@ TritonModel::Create(
7575
}
7676

7777
// Localize the content of the model repository corresponding to
78-
// 'model_path'. This model holds a handle to the localized content
79-
// so that it persists as long as the model is loaded.
78+
// 'model_path' and model's version. This model holds a handle to
79+
// the localized content so that it persists as long as the model is loaded.
8080
std::shared_ptr<LocalizedPath> localized_model_dir;
81-
RETURN_IF_ERROR(LocalizePath(model_path, &localized_model_dir));
81+
RETURN_IF_ERROR(
82+
LocalizePath(model_path, std::to_string(version), &localized_model_dir));
8283

8384
// Localize paths in backend model config
8485
// [FIXME] Remove once a more permanent solution is implemented (DLIS-4211)

‎src/filesystem/api.cc

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -558,11 +558,13 @@ ReadTextProto(const std::string& path, google::protobuf::Message* msg)
558558
}
559559

560560
Status
561-
LocalizePath(const std::string& path, std::shared_ptr<LocalizedPath>* localized)
561+
LocalizePath(
562+
const std::string& path, const std::string& fetch_subdir,
563+
std::shared_ptr<LocalizedPath>* localized)
562564
{
563565
std::shared_ptr<FileSystem> fs;
564566
RETURN_IF_ERROR(fsm_.GetFileSystem(path, fs));
565-
return fs->LocalizePath(path, localized);
567+
return fs->LocalizePath(path, fetch_subdir, localized);
566568
}
567569

568570
Status

‎src/filesystem/api.h

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -145,11 +145,15 @@ Status ReadTextFile(const std::string& path, std::string* contents);
145145

146146
/// Create an object representing a local copy of a path.
147147
/// \param path The path of the directory or file.
148+
/// \param fetch_subdir If specified, will only download provided
149+
/// sub directory, otherwise all subdirectories will be downloaded.
150+
/// Does not affect files individual files, located under `path`.
148151
/// \param localized Returns the LocalizedPath object
149152
/// representing the local copy of the path.
150153
/// \return Error status
151154
Status LocalizePath(
152-
const std::string& path, std::shared_ptr<LocalizedPath>* localized);
155+
const std::string& path, const std::string& fetch_subdir,
156+
std::shared_ptr<LocalizedPath>* localized);
153157

154158
/// Write a string to a file.
155159
/// \param path The path of the file.

‎src/filesystem/implementations/as.h

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -86,7 +86,7 @@ class ASFileSystem : public FileSystem {
8686
const std::string& path, std::set<std::string>* files) override;
8787
Status ReadTextFile(const std::string& path, std::string* contents) override;
8888
Status LocalizePath(
89-
const std::string& path,
89+
const std::string& path, const std::string& fetch_subdir,
9090
std::shared_ptr<LocalizedPath>* localized) override;
9191
Status WriteTextFile(
9292
const std::string& path, const std::string& contents) override;
@@ -424,7 +424,8 @@ ASFileSystem::DownloadFolder(
424424

425425
Status
426426
ASFileSystem::LocalizePath(
427-
const std::string& path, std::shared_ptr<LocalizedPath>* localized)
427+
const std::string& path, const std::string& fetch_subdir,
428+
std::shared_ptr<LocalizedPath>* localized)
428429
{
429430
bool exists;
430431
RETURN_IF_ERROR(FileExists(path, &exists));

‎src/filesystem/implementations/common.h

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -83,7 +83,8 @@ class FileSystem {
8383
virtual Status ReadTextFile(
8484
const std::string& path, std::string* contents) = 0;
8585
virtual Status LocalizePath(
86-
const std::string& path, std::shared_ptr<LocalizedPath>* localized) = 0;
86+
const std::string& path, const std::string& fetch_subdir,
87+
std::shared_ptr<LocalizedPath>* localized) = 0;
8788
virtual Status WriteTextFile(
8889
const std::string& path, const std::string& contents) = 0;
8990
virtual Status WriteBinaryFile(

‎src/filesystem/implementations/gcs.h

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -75,7 +75,7 @@ class GCSFileSystem : public FileSystem {
7575
const std::string& path, std::set<std::string>* files) override;
7676
Status ReadTextFile(const std::string& path, std::string* contents) override;
7777
Status LocalizePath(
78-
const std::string& path,
78+
const std::string& path, const std::string& fetch_subdir,
7979
std::shared_ptr<LocalizedPath>* localized) override;
8080
Status WriteTextFile(
8181
const std::string& path, const std::string& contents) override;
@@ -363,7 +363,8 @@ GCSFileSystem::ReadTextFile(const std::string& path, std::string* contents)
363363

364364
Status
365365
GCSFileSystem::LocalizePath(
366-
const std::string& path, std::shared_ptr<LocalizedPath>* localized)
366+
const std::string& path, const std::string& fetch_subdir,
367+
std::shared_ptr<LocalizedPath>* localized)
367368
{
368369
bool exists;
369370
RETURN_IF_ERROR(FileExists(path, &exists));

‎src/filesystem/implementations/local.h

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,7 @@ class LocalFileSystem : public FileSystem {
4949
const std::string& path, std::set<std::string>* files) override;
5050
Status ReadTextFile(const std::string& path, std::string* contents) override;
5151
Status LocalizePath(
52-
const std::string& path,
52+
const std::string& path, const std::string& fetch_subdir,
5353
std::shared_ptr<LocalizedPath>* localized) override;
5454
Status WriteTextFile(
5555
const std::string& path, const std::string& contents) override;
@@ -204,7 +204,8 @@ LocalFileSystem::ReadTextFile(const std::string& path, std::string* contents)
204204

205205
Status
206206
LocalFileSystem::LocalizePath(
207-
const std::string& path, std::shared_ptr<LocalizedPath>* localized)
207+
const std::string& path, const std::string& fetch_subdir,
208+
std::shared_ptr<LocalizedPath>* localized)
208209
{
209210
// For local file system we don't actually need to download the
210211
// directory or file. We use it in place.

‎src/filesystem/implementations/s3.h

Lines changed: 27 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -151,7 +151,7 @@ class S3FileSystem : public FileSystem {
151151
const std::string& path, std::set<std::string>* files) override;
152152
Status ReadTextFile(const std::string& path, std::string* contents) override;
153153
Status LocalizePath(
154-
const std::string& path,
154+
const std::string& path, const std::string& fetch_subdir,
155155
std::shared_ptr<LocalizedPath>* localized) override;
156156
Status WriteTextFile(
157157
const std::string& path, const std::string& contents) override;
@@ -628,7 +628,8 @@ S3FileSystem::ReadTextFile(const std::string& path, std::string* contents)
628628

629629
Status
630630
S3FileSystem::LocalizePath(
631-
const std::string& path, std::shared_ptr<LocalizedPath>* localized)
631+
const std::string& path, const std::string& fetch_subdir,
632+
std::shared_ptr<LocalizedPath>* localized)
632633
{
633634
// Check if the directory or file exists
634635
bool exists;
@@ -693,28 +694,34 @@ S3FileSystem::LocalizePath(
693694
: JoinPath({(*localized)->Path(), s3_removed_path});
694695
bool is_subdir;
695696
RETURN_IF_ERROR(IsDirectory(s3_fpath, &is_subdir));
697+
bool copy_subdir =
698+
!fetch_subdir.empty()
699+
? s3_fpath == JoinPath({effective_path, fetch_subdir})
700+
: true;
696701
if (is_subdir) {
697-
// Create local mirror of sub-directories
702+
if (copy_subdir) {
703+
// Create local mirror of sub-directories
698704
#ifdef _WIN32
699-
int status = mkdir(const_cast<char*>(local_fpath.c_str()));
705+
int status = mkdir(const_cast<char*>(local_fpath.c_str()));
700706
#else
701-
int status = mkdir(
702-
const_cast<char*>(local_fpath.c_str()),
703-
S_IRUSR | S_IWUSR | S_IXUSR);
707+
int status = mkdir(
708+
const_cast<char*>(local_fpath.c_str()),
709+
S_IRUSR | S_IWUSR | S_IXUSR);
704710
#endif
705-
if (status == -1) {
706-
return Status(
707-
Status::Code::INTERNAL,
708-
"Failed to create local folder: " + local_fpath +
709-
", errno:" + strerror(errno));
710-
}
711-
712-
// Add sub-directories and deeper files to contents
713-
std::set<std::string> subdir_contents;
714-
RETURN_IF_ERROR(GetDirectoryContents(s3_fpath, &subdir_contents));
715-
for (auto itr = subdir_contents.begin(); itr != subdir_contents.end();
716-
++itr) {
717-
contents.insert(JoinPath({s3_fpath, *itr}));
711+
if (status == -1) {
712+
return Status(
713+
Status::Code::INTERNAL,
714+
"Failed to create local folder: " + local_fpath +
715+
", errno:" + strerror(errno));
716+
}
717+
718+
// Add sub-directories and deeper files to contents
719+
std::set<std::string> subdir_contents;
720+
RETURN_IF_ERROR(GetDirectoryContents(s3_fpath, &subdir_contents));
721+
for (auto itr = subdir_contents.begin(); itr != subdir_contents.end();
722+
++itr) {
723+
contents.insert(JoinPath({s3_fpath, *itr}));
724+
}
718725
}
719726
} else {
720727
// Create local copy of file

‎src/model_config_utils.cc

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -918,8 +918,8 @@ LocalizePythonBackendExecutionEnvironmentPath(
918918
model_path_slash) {
919919
// Localize the file
920920
std::shared_ptr<LocalizedPath> localized_exec_env_path;
921-
RETURN_IF_ERROR(
922-
LocalizePath(abs_exec_env_path, &localized_exec_env_path));
921+
RETURN_IF_ERROR(LocalizePath(
922+
abs_exec_env_path, "" /*fetch_subdir*/, &localized_exec_env_path));
923923
// Persist the localized temporary path
924924
(*localized_model_dir)
925925
->other_localized_path.push_back(localized_exec_env_path);

0 commit comments

Comments
 (0)