[SageMaker Galactus developer experience] use python backend when sch…

…ema is customized
deepjavalibrary · Nov 7, 2023 · 5221daf · 5221daf
1 parent 407aea4
commit 5221daf
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 4 deletions.
diff --git a/wlm/src/main/java/ai/djl/serving/wlm/SageMakerUtils.java b/wlm/src/main/java/ai/djl/serving/wlm/SageMakerUtils.java
@@ -44,7 +44,6 @@ public final class SageMakerUtils {
     private static final String TASK = "Task";
     private static final String VERSION = "Version";
     private static final String XGBOOST_MODEL_TYPE = "XGBoostModel";
-    private static final String YAML_FORMAT = ".yaml";
 
     private SageMakerUtils() {}
 
@@ -110,8 +109,7 @@ private static boolean validateMetaData(Map<String, Object> metaDataMap) {
 
     private static boolean hasCustomizedSchema(Map<String, Object> metaDataMap)
             throws ModelException {
-        if (metaDataMap.containsKey(SCHEMA)
-                && !metaDataMap.get(SCHEMA).toString().contains(YAML_FORMAT)) {
+        if (metaDataMap.containsKey(SCHEMA)) {
             if (!metaDataMap.containsKey(SCHEMA_HMAC)) {
                 throw new ModelException(
                         "Invalid SageMaker Model format due to SchemaHMAC is not found");

diff --git a/wlm/src/test/java/ai/djl/serving/wlm/ModelInfoTest.java b/wlm/src/test/java/ai/djl/serving/wlm/ModelInfoTest.java
@@ -271,7 +271,8 @@ public void testInferLMIEngine() throws IOException, ModelException {
     public void testInferSageMakerEngine() throws IOException, ModelException {
         ModelInfo<Input, Output> model = new ModelInfo<>("src/test/resources/sagemaker/xgb_model");
         model.initialize();
-        assertEquals(model.getEngineName(), "XGBoost");
+        assertEquals(model.getEngineName(), "Python");
+        assertEquals(model.prop.getProperty("option.entryPoint"), "djl_python.sagemaker");
 
         // test case for valid pytorch model with customized schema and inference spec
         model = new ModelInfo<>("src/test/resources/sagemaker/pytorch_model");