diff --git a/serving/docker/partition/partition.py b/serving/docker/partition/partition.py index ecaeed1be..35f38b5f5 100644 --- a/serving/docker/partition/partition.py +++ b/serving/docker/partition/partition.py @@ -42,7 +42,7 @@ def __init__(self, props_manager): self.download_model_from_s3() def download_model_from_s3(self): - model_id = self.properties.get("model_id") + model_id = self.properties.get("option.model_id") if not model_id or not model_id.startswith("s3://"): return @@ -71,7 +71,7 @@ def download_model_from_s3(self): if not glob.glob(os.path.join(download_dir, '*')): raise Exception('Model download from s3url failed') - self.properties['model_id'] = download_dir + self.properties['option.model_id'] = download_dir def install_requirements_file(self): req_file_dir = self.properties_manager.properties_dir @@ -106,13 +106,13 @@ def set_environmental_vars(self): def download_config_from_hf(self): # checks if model_id is a path - if glob.glob(self.properties['model_id']): - return self.properties['model_id'] + if glob.glob(self.properties['option.model_id']): + return self.properties['option.model_id'] download_dir = os.environ.get("SERVING_DOWNLOAD_DIR", '/tmp/download/model/') - model_name = self.properties['model_id'] + model_name = self.properties['option.model_id'] downloaded_dir = snapshot_download( repo_id=model_name, cache_dir=download_dir, @@ -122,7 +122,7 @@ def download_config_from_hf(self): def copy_config_files(self): model_dir = self.properties['model_dir'] - if 'model_id' in self.properties: + if 'option.model_id' in self.properties: model_dir = self.download_config_from_hf() config_files = [] @@ -130,14 +130,16 @@ def copy_config_files(self): config_files += glob.glob(os.path.join(model_dir, pattern)) for file in config_files: - shutil.copy(file, dst=self.properties['save_mp_checkpoint_path']) + shutil.copy(file, + dst=self.properties['option.save_mp_checkpoint_path']) def upload_checkpoints_to_s3(self): if 'upload_checkpoints_s3url' not in self.properties: return s3url = self.properties['upload_checkpoints_s3url'] - saved_checkpoints_dir = self.properties["save_mp_checkpoint_path"] + saved_checkpoints_dir = self.properties[ + "option.save_mp_checkpoint_path"] if not saved_checkpoints_dir.endswith('/'): saved_checkpoints_dir = saved_checkpoints_dir + '/' @@ -154,7 +156,7 @@ def upload_checkpoints_to_s3(self): commands = ["aws", "s3", "sync", saved_checkpoints_dir, s3url] subprocess.run(commands) - shutil.rmtree(self.properties["save_mp_checkpoint_path"]) + shutil.rmtree(self.properties["option.save_mp_checkpoint_path"]) def cleanup(self): """ @@ -184,16 +186,18 @@ def run_partition(self): def load_the_generated_checkpoints(self): if self.properties['engine'] == 'DeepSpeed': - saved_checkpoints_dir = self.properties["save_mp_checkpoint_path"] + saved_checkpoints_dir = self.properties[ + "option.save_mp_checkpoint_path"] properties = utils.load_properties(saved_checkpoints_dir) properties['model_dir'] = saved_checkpoints_dir - properties['entryPoint'] = self.properties['entryPoint'] + properties['option.entryPoint'] = self.properties[ + 'option.entryPoint'] properties['partition_handler'] = 'handle' entry_point_file = None - if properties['entryPoint'] == 'model.py': + if properties['option.entryPoint'] == 'model.py': entry_point_file = os.path.join( - self.properties['properties_dir'], 'model.py') + self.properties_manager.properties_dir, 'model.py') shutil.copy(entry_point_file, saved_checkpoints_dir) commands = get_partition_cmd(True, properties) diff --git a/serving/docker/partition/properties_manager.py b/serving/docker/partition/properties_manager.py index 927da4290..e5ee85810 100644 --- a/serving/docker/partition/properties_manager.py +++ b/serving/docker/partition/properties_manager.py @@ -20,22 +20,17 @@ from utils import is_engine_mpi_mode, get_engine_configs, get_download_dir, load_properties EXCLUDE_PROPERTIES = [ - 'option.model_id', 'save_mp_checkpoint_path', 'model_dir', + 'option.model_id', 'option.save_mp_checkpoint_path', 'model_dir', 'upload_checkpoints_s3url', 'properties_dir' ] -PARTITION_SUPPORTED_ENGINES = ['DeepSpeed', 'FasterTransformer'] - -CHUNK_SIZE = 4096 # 4MB chunk size - class PropertiesManager(object): def __init__(self, args): self.entry_point_url = None - self.properties = {} self.properties_dir = args.properties_dir - load_properties(self.properties_dir) + self.properties = load_properties(self.properties_dir) if args.model_id: self.properties['option.model_id'] = args.model_id @@ -43,7 +38,7 @@ def __init__(self, args): self.properties['engine'] = args.engine if args.save_mp_checkpoint_path: self.properties[ - 'save_mp_checkpoint_path'] = args.save_mp_checkpoint_path + 'option.save_mp_checkpoint_path'] = args.save_mp_checkpoint_path if args.tensor_parallel_degree: self.properties[ 'option.tensor_parallel_degree'] = args.tensor_parallel_degree @@ -86,7 +81,7 @@ def validate_and_correct_checkpoints_json(self): """ if self.properties.get('engine') == 'DeepSpeed': config_file = os.path.join( - self.properties['save_mp_checkpoint_path'], + self.properties['option.save_mp_checkpoint_path'], 'ds_inference_config.json') if not os.path.exists(config_file): raise ValueError("Checkpoints json file was not generated." @@ -103,17 +98,17 @@ def validate_and_correct_checkpoints_json(self): json.dump(configs, f) def generate_properties_file(self): - checkpoint_path = self.properties.get('save_mp_checkpoint_path') + checkpoint_path = self.properties.get('option.save_mp_checkpoint_path') configs = get_engine_configs(self.properties) for key, value in self.properties.items(): if key not in EXCLUDE_PROPERTIES: - if key == "entryPoint": - entry_point = self.properties.get("entryPoint") + if key == "option.entryPoint": + entry_point = self.properties.get("option.entryPoint") if entry_point == "model.py": continue elif self.entry_point_url: - configs["entryPoint"] = self.entry_point_url + configs["option.entryPoint"] = self.entry_point_url else: configs[key] = value @@ -135,14 +130,14 @@ def validate_tp_degree(self): ) def set_and_validate_entry_point(self): - entry_point = self.properties.get('entryPoint') + entry_point = self.properties.get('option.entryPoint') if entry_point is None: entry_point = os.environ.get("DJL_ENTRY_POINT") if entry_point is None: entry_point_file = glob.glob( os.path.join(self.properties_dir, 'model.py')) if entry_point_file: - self.properties['entryPoint'] = 'model.py' + self.properties['option.entryPoint'] = 'model.py' else: engine = self.properties.get('engine') if engine == "DeepSpeed": @@ -151,7 +146,7 @@ def set_and_validate_entry_point(self): entry_point = "djl_python.fastertransformer" else: raise ValueError("Please specify engine") - self.properties['entryPoint'] = entry_point + self.properties['option.entryPoint'] = entry_point elif entry_point.lower().startswith('http'): logging.info(f'Downloading entrypoint file.') self.entry_point_url = entry_point @@ -161,16 +156,17 @@ def set_and_validate_entry_point(self): with requests.get(entry_point) as r: with open(model_file, 'wb') as f: f.write(r.content) - self.properties['entryPoint'] = model_file + self.properties['option.entryPoint'] = model_file logging.info(f'Entrypoint file downloaded successfully') def set_and_validate_save_mp_checkpoint_path(self): save_mp_checkpoint_path = self.properties.get( - "save_mp_checkpoint_path") + "option.save_mp_checkpoint_path") if not save_mp_checkpoint_path: raise ValueError("Please specify save_mp_checkpoint_path") if save_mp_checkpoint_path.startswith("s3://"): self.properties[ "upload_checkpoints_s3url"] = save_mp_checkpoint_path - self.properties["save_mp_checkpoint_path"] = get_download_dir( - self.properties_dir, "partition-model") + self.properties[ + "option.save_mp_checkpoint_path"] = get_download_dir( + self.properties_dir, "partition-model") diff --git a/serving/docker/partition/utils.py b/serving/docker/partition/utils.py index ccb3ea6ef..faf6221a1 100644 --- a/serving/docker/partition/utils.py +++ b/serving/docker/partition/utils.py @@ -90,5 +90,5 @@ def load_properties(properties_dir): if line.startswith("#") or not line.strip(): continue key, value = line.strip().split('=', 1) - properties[key.split(".", 1)[-1]] = value + properties[key] = value return properties