From 1d194a9fe20b1394a1f415de396e36536da1f136 Mon Sep 17 00:00:00 2001 From: Somasundaram Date: Mon, 27 Jan 2025 17:25:29 -0800 Subject: [PATCH] [python] remove deepspeed related AOT code --- serving/docker/partition/partition.py | 33 ------------------- .../docker/partition/properties_manager.py | 33 +------------------ serving/docker/partition/utils.py | 12 +------ 3 files changed, 2 insertions(+), 76 deletions(-) diff --git a/serving/docker/partition/partition.py b/serving/docker/partition/partition.py index 768fbb364..96ca4c121 100644 --- a/serving/docker/partition/partition.py +++ b/serving/docker/partition/partition.py @@ -20,7 +20,6 @@ from pathlib import Path -import utils from properties_manager import PropertiesManager from huggingface_hub import snapshot_download from datasets import load_dataset @@ -196,12 +195,10 @@ def run_partition(self) -> str: logging.info(proc) if proc.returncode == 0: logging.info("Partitioning done.") - self.properties_manager.validate_and_correct_checkpoints_json() self.properties_manager.generate_properties_file() if not self.properties_manager.skip_copy: logging.info("Copying config files...") self.copy_config_files() - self.load_the_generated_checkpoints() self.upload_checkpoints_to_s3() self.cleanup() return partition_stdout @@ -212,36 +209,6 @@ def run_partition(self) -> str: f"Partitioning exited with return code: {proc.returncode}. Details: {partition_stderr}" ) - def load_the_generated_checkpoints(self): - if self.properties['engine'] == 'DeepSpeed': - saved_checkpoints_dir = self.properties[ - "option.save_mp_checkpoint_path"] - properties = utils.load_properties(saved_checkpoints_dir) - if not self.properties_manager.skip_copy: - properties['model_dir'] = saved_checkpoints_dir - properties['option.entryPoint'] = self.properties[ - 'option.entryPoint'] - properties['partition_handler'] = 'handle' - - entry_point_file = None - if properties['option.entryPoint'] == 'model.py': - entry_point_file = os.path.join( - self.properties_manager.properties_dir, 'model.py') - shutil.copy(entry_point_file, saved_checkpoints_dir) - - commands = get_partition_cmd(True, properties) - self.set_environmental_vars() - result = subprocess.run(commands) - logging.info(result) - if result.returncode == 0: - logging.info( - "Successfully loaded the partitioned checkpoints.") - else: - raise Exception("DeepSpeed does not support partitioning. " - "Please use a different engine") - if entry_point_file: - os.remove(os.path.join(saved_checkpoints_dir, 'model.py')) - def run_quantization(self): quant_method = self.properties['option.quantize'] if quant_method == 'awq': diff --git a/serving/docker/partition/properties_manager.py b/serving/docker/partition/properties_manager.py index 70d7128e9..1bccbb289 100644 --- a/serving/docker/partition/properties_manager.py +++ b/serving/docker/partition/properties_manager.py @@ -12,12 +12,11 @@ import logging import os import glob -import json import torch import requests # Properties to exclude while generating serving.properties -from utils import (is_engine_mpi_mode, get_engine_configs, get_download_dir, +from utils import (is_engine_mpi_mode, get_download_dir, load_properties, update_kwargs_with_env_vars) EXCLUDE_PROPERTIES = [ @@ -88,34 +87,6 @@ def set_and_validate_model_dir(self): f'No .bin or .safetensors files found in the dir: {self.properties_dir}' '\nPlease specify the model_dir or model_id') - def validate_and_correct_checkpoints_json(self): - """ - Removes base_dir from ds_inference_checkpoints.json file. - - DeepSpeed writes base_dir directory, which is the path of checkpoints saved to the file. - Removing the base_dir since the user's deployment environment could be different from partition environment. - User can specify base_dir argument in deepspeed.init_inference while using this file. - - :return: - """ - if self.properties.get('engine') == 'DeepSpeed': - config_file = os.path.join( - 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." - "Partition was not successful.") - - with open(config_file) as f: - configs = json.load(f) - - if not configs.get('base_dir'): - return - - configs.pop('base_dir') - with open(config_file, "w") as f: - json.dump(configs, f) - def generate_properties_file(self): checkpoint_path = self.properties.get('option.save_mp_checkpoint_path') configs = get_engine_configs(self.properties) @@ -172,8 +143,6 @@ def set_and_validate_entry_point(self): pass elif engine is None: raise ValueError("Please specify engine") - elif engine.lower() == "deepspeed": - entry_point = "djl_python.deepspeed" elif engine.lower() == "python": entry_point = "djl_python.transformers_neuronx" else: diff --git a/serving/docker/partition/utils.py b/serving/docker/partition/utils.py index 802523085..6d2dcb48c 100644 --- a/serving/docker/partition/utils.py +++ b/serving/docker/partition/utils.py @@ -42,16 +42,6 @@ def get_partition_cmd(is_mpi_mode, properties): ] -def get_engine_configs(properties): - engine = properties.get('engine') - configs = {} - if engine == 'DeepSpeed': - configs['option.checkpoint'] = 'ds_inference_config.json' - configs['option.parallel_loading'] = True - - return configs - - def extract_python_jar(target_dir): os.makedirs(target_dir, exist_ok=True) jar_files = glob.glob('/usr/local/djl-serving-*/lib/python-*.jar') @@ -72,7 +62,7 @@ def get_djl_version_from_lib(): def is_engine_mpi_mode(engine): - if engine == 'DeepSpeed': + if engine == 'MPI': return True else: return False