Skip to content

Commit

Permalink
workaround solution for SageMaker (#2405)
Browse files Browse the repository at this point in the history
  • Loading branch information
lxning authored Jun 13, 2023
1 parent 253f8a3 commit 3917407
Showing 1 changed file with 15 additions and 2 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -80,15 +80,28 @@ public static ModelArchive downloadModel(
}
}

if (new File(url).isDirectory()) {
File directory = new File(url);
if (directory.isDirectory()) {
// handle the case that the input url is a directory.
// the input of url is "/xxx/model_store/modelXXX" or
// "xxxx/yyyyy/modelXXX".
return load(url, new File(url), false);
File[] fileList = directory.listFiles();
if (fileList.length == 1 && fileList[0].isDirectory()) {
// handle the case that a model tgz file
// has root dir after decompression on SageMaker
return load(url, fileList[0], false);
}
return load(url, directory, false);
} else if (modelLocation.exists()) {
// handle the case that "/xxx/model_store/modelXXX" is directory.
// the input of url is modelXXX when torchserve is started
// with snapshot or with parameter --models modelXXX
File[] fileList = modelLocation.listFiles();
if (fileList.length == 1 && fileList[0].isDirectory()) {
// handle the case that a model tgz file
// has root dir after decompression on SageMaker
return load(url, fileList[0], false);
}
return load(url, modelLocation, false);
}

Expand Down

0 comments on commit 3917407

Please sign in to comment.