[neo] Get SM Fast Model Loader model dir name from package (#2541)

deepjavalibrary · Nov 11, 2024 · 3b91858 · 3b91858
1 parent 8481f56
commit 3b91858
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/serving/docker/partition/sm_neo_shard.py b/serving/docker/partition/sm_neo_shard.py
@@ -30,7 +30,6 @@
 from lmi_dist.arg_utils import VllmEngineArgs
 from lmi_dist.comms import comms
 
-MODEL_DIR_NAME = "sagemaker-fast-model-loader"
 CHUNK_MB = 8
 
 
@@ -107,7 +106,7 @@ def shard_lmi_dist_model(self, input_dir: str, output_dir: str,
         )
         engine = engine_from_args(engine_args)
 
-        model_dir = os.path.join(output_dir, MODEL_DIR_NAME)
+        model_dir = os.path.join(output_dir, sm_fml.MODEL_DIR_NAME)
         os.makedirs(model_dir, exist_ok=True)
 
         config_for_current_rank = engine.model_runner.vllm_worker.save_chunked_shard(