From 8cbf7b933a6869e6c117eb1bc2d22aaf7efababa Mon Sep 17 00:00:00 2001 From: LuLu Zuo Date: Thu, 30 Nov 2023 15:53:25 +0800 Subject: [PATCH 01/32] Optimize memory computation --- src/promptflow/promptflow/executor/_errors.py | 6 ++ .../executor/_line_execution_process_pool.py | 58 ++++++++++++++----- .../unittests/executor/test_flow_executor.py | 49 ++++++++++++---- 3 files changed, 86 insertions(+), 27 deletions(-) diff --git a/src/promptflow/promptflow/executor/_errors.py b/src/promptflow/promptflow/executor/_errors.py index 9ec5b3d10c2..321627d3522 100644 --- a/src/promptflow/promptflow/executor/_errors.py +++ b/src/promptflow/promptflow/executor/_errors.py @@ -175,6 +175,12 @@ def __init__(self, line_number, timeout): ) +class InvalidMemoryUsageFactor(UserErrorException): + """Exception raised when node reference not found or unsupported""" + + pass + + class EmptyLLMApiMapping(UserErrorException): """Exception raised when connection_type_to_api_mapping is empty and llm node provider can't be inferred""" diff --git a/src/promptflow/promptflow/executor/_line_execution_process_pool.py b/src/promptflow/promptflow/executor/_line_execution_process_pool.py index 1da653b4a49..dbc18b55aec 100644 --- a/src/promptflow/promptflow/executor/_line_execution_process_pool.py +++ b/src/promptflow/promptflow/executor/_line_execution_process_pool.py @@ -27,7 +27,7 @@ from promptflow.contracts.run_info import RunInfo as NodeRunInfo from promptflow.contracts.run_info import Status from promptflow.exceptions import ErrorTarget, PromptflowException -from promptflow.executor._errors import LineExecutionTimeoutError +from promptflow.executor._errors import LineExecutionTimeoutError, InvalidMemoryUsageFactor from promptflow.executor._result import LineResult from promptflow.executor.flow_executor import DEFAULT_CONCURRENCY_BULK, FlowExecutor from promptflow.storage import AbstractRunStorage @@ -197,12 +197,19 @@ def __enter__(self): # Starting a new process in non-fork mode requires to allocate memory. Determine the maximum number of processes # based on available memory to avoid memory bursting. if not self._use_fork: - available_max_worker_count = get_available_max_worker_count() + memory_usage_factor = os.environ.get("PF_BATCH_Memory_Usage_Factor") + available_max_worker_count = get_available_max_worker_count(memory_usage_factor) self._n_process = min(self._worker_count, self._nlines, available_max_worker_count) - bulk_logger.info(f"Not using fork, process count: {self._n_process}") + bulk_logger.info( + f"Not using fork, the number of processes is determined by calculating the difference between the" + f"current system available memory and the total memory usage factor set by the environment" + f"variable multiplied by the total memory, and then dividing this difference by the memory of process." + f"If not set memory usage factor, the default is 1, process count: {self._n_process}") else: self._n_process = min(self._worker_count, self._nlines) - bulk_logger.info(f"Using fork, process count: {self._n_process}") + bulk_logger.info( + f"Using fork, the number of processes is determined by the lesser of the worker_count set by the" + f"environment variable configuration and the row count, If not process count: {self._n_process}") pool = ThreadPool(self._n_process, initializer=set_context, initargs=(contextvars.copy_context(),)) self._pool = pool @@ -537,22 +544,21 @@ def create_executor_legacy(*, flow, connections, loaded_tools, cache_manager, st ) -def get_available_max_worker_count(): +def get_available_max_worker_count(memory_usage_factor=None): + if memory_usage_factor is None: + memory_usage_factor = 1 + memory_usage_factor = convert_to_decimal(memory_usage_factor) pid = os.getpid() mem_info = psutil.virtual_memory() total_memory = mem_info.total / (1024 * 1024) # in MB - total_memory_in_use = mem_info.used / (1024 * 1024) # in MB - available_memory = mem_info.available / (1024 * 1024) # in MB + sys_available_memory = mem_info.available / (1024 * 1024) # in MB process = psutil.Process(pid) process_memory_info = process.memory_info() process_memory = process_memory_info.rss / (1024 * 1024) # in MB # To ensure system stability, reserve memory for system usage. - available_max_worker_count = math.floor((available_memory - 0.3 * total_memory) / process_memory) + available_memory = (sys_available_memory - (1-memory_usage_factor) * total_memory) + available_max_worker_count = math.floor(available_memory / process_memory) if available_max_worker_count < 1: - # For the case of vector db, at most 1/3 of the memory will be used, which is 33% of the memory - # In this scenario, the "available_max_worker_count" may be 0, which will cause an error - # "Number of processes must be at least 1" when creating ThreadPool - # So set "available_max_worker_count" to 1 if it's less than 1 # TODO: For the case of vector db, Optimize execution logic # 1. Let the main process not consume memory because it does not actually invoke # 2. When the degree of parallelism is 1, main process executes the task directly and not @@ -560,9 +566,9 @@ def get_available_max_worker_count(): bulk_logger.warning(f"Available max worker count {available_max_worker_count} is less than 1, set it to 1.") available_max_worker_count = 1 bulk_logger.info( - f"""Process {pid} uses {process_memory}, - total memory {total_memory}, total memory in use: {total_memory_in_use}, - available memory: {available_memory}, available max worker count: {available_max_worker_count}""" + f"""Process {pid} current available memory is {process_memory}, + memory consumption of current process is {available_memory}, + worker count is set to {available_memory}/{process_memory} = {available_max_worker_count}""" ) return available_max_worker_count @@ -575,3 +581,25 @@ def get_multiprocessing_context(multiprocessing_start_method=None): else: context = multiprocessing.get_context() return context + + +def convert_to_decimal(memory_usage_factor): + if isinstance(memory_usage_factor, (int, float)): + return memory_usage_factor + elif memory_usage_factor.endswith('%'): + # Remove the '%' sign and convert to float + number = float(memory_usage_factor.rstrip('%')) + # Convert the percentage to a decimal + return number / 100 + else: + try: + # Try to convert the string to a float + return float(memory_usage_factor) + except ValueError: + # If the string cannot be converted to a float, raise an error + raise InvalidMemoryUsageFactor( + message_format=( + f"Cannot convert string to number: {memory_usage_factor}" + ), + memory_usage_factor=memory_usage_factor + ) diff --git a/src/promptflow/tests/executor/unittests/executor/test_flow_executor.py b/src/promptflow/tests/executor/unittests/executor/test_flow_executor.py index 268a37dd20d..b0f5e167e12 100644 --- a/src/promptflow/tests/executor/unittests/executor/test_flow_executor.py +++ b/src/promptflow/tests/executor/unittests/executor/test_flow_executor.py @@ -6,6 +6,8 @@ from promptflow.contracts.flow import FlowInputDefinition from promptflow.contracts.tool import ValueType from promptflow.executor._line_execution_process_pool import get_available_max_worker_count +from promptflow.executor._errors import InvalidMemoryUsageFactor + from promptflow.executor.flow_executor import ( FlowExecutor, _ensure_node_result_is_serializable, @@ -182,15 +184,34 @@ def test_non_streaming_tool_should_not_be_affected(self): class TestGetAvailableMaxWorkerCount: @pytest.mark.parametrize( - "total_memory, available_memory, process_memory, expected_max_worker_count, actual_calculate_worker_count", + ( + "total_memory", + "available_memory", + "process_memory", + "memory_usage_factor", + "expected_max_worker_count", + "actual_calculate_worker_count", + "expect_exception" + ), [ - (1024.0, 128.0, 64.0, 1, -3), # available_memory - 0.3 * total_memory < 0 - (1024.0, 307.20, 64.0, 1, 0), # available_memory - 0.3 * total_memory = 0 - (1024.0, 768.0, 64.0, 7, 7), # available_memory - 0.3 * total_memory > 0 + (1024.0, 128.0, 64.0, 0, 1, -14, False), # available_memory - memory_usage_factor * total_memory < 0 + (1024.0, 512.0, 64.0, 0.5, 1, 0, False), # available_memory - memory_usage_factor * total_memory = 0 + (1024.0, 768.0, 64.0, 0.7, 7, 7, False), # available_memory - memory_usage_factor * total_memory > 0 + (1024.0, 768.0, 64.0, 0.8, 8, 8, False), # memory_usage_factor is None + (1024.0, 512.0, 64.0, "50%", 1, 0, False), # The environment variable is configured as a percentage + (1024.0, 512.0, 64.0, "0.5", 1, 0, False), # The environment variable is configured as a string + (1024.0, 512.0, 64.0, "0.5a", 1, 0, True) ], ) def test_get_available_max_worker_count( - self, total_memory, available_memory, process_memory, expected_max_worker_count, actual_calculate_worker_count + self, + total_memory, + available_memory, + process_memory, + memory_usage_factor, + expected_max_worker_count, + actual_calculate_worker_count, + expect_exception ): with patch("psutil.virtual_memory") as mock_mem: mock_mem.return_value.total = total_memory * 1024 * 1024 @@ -199,10 +220,14 @@ def test_get_available_max_worker_count( mock_process.return_value.memory_info.return_value.rss = process_memory * 1024 * 1024 with patch("promptflow.executor._line_execution_process_pool.bulk_logger") as mock_logger: mock_logger.warning.return_value = None - max_worker_count = get_available_max_worker_count() - assert max_worker_count == expected_max_worker_count - if actual_calculate_worker_count < 1: - mock_logger.warning.assert_called_with( - f"Available max worker count {actual_calculate_worker_count} is less than 1, " - "set it to 1." - ) + if expect_exception: + with pytest.raises(InvalidMemoryUsageFactor): + get_available_max_worker_count(memory_usage_factor) + else: + max_worker_count = get_available_max_worker_count(memory_usage_factor) + assert max_worker_count == expected_max_worker_count + if actual_calculate_worker_count < 1: + mock_logger.warning.assert_called_with( + f"Available max worker count {actual_calculate_worker_count} is less than 1, " + "set it to 1." + ) From 3aae7dfb5f116136389cea719b3aa9adc6c18628 Mon Sep 17 00:00:00 2001 From: LuLu Zuo Date: Thu, 30 Nov 2023 15:55:15 +0800 Subject: [PATCH 02/32] Optimize memory computation --- .../promptflow/executor/_line_execution_process_pool.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/promptflow/promptflow/executor/_line_execution_process_pool.py b/src/promptflow/promptflow/executor/_line_execution_process_pool.py index dbc18b55aec..3c791a33d9a 100644 --- a/src/promptflow/promptflow/executor/_line_execution_process_pool.py +++ b/src/promptflow/promptflow/executor/_line_execution_process_pool.py @@ -209,7 +209,8 @@ def __enter__(self): self._n_process = min(self._worker_count, self._nlines) bulk_logger.info( f"Using fork, the number of processes is determined by the lesser of the worker_count set by the" - f"environment variable configuration and the row count, If not process count: {self._n_process}") + f"environment variable configuration and the row count, If not set the worker_count, the default" + f"is 16, process count: {self._n_process}") pool = ThreadPool(self._n_process, initializer=set_context, initargs=(contextvars.copy_context(),)) self._pool = pool From 45f9c28dc1f3220e110b90e616dbf2d3695f7411 Mon Sep 17 00:00:00 2001 From: LuLu Zuo Date: Thu, 30 Nov 2023 18:58:48 +0800 Subject: [PATCH 03/32] Optimize memory computation --- src/promptflow/promptflow/executor/_errors.py | 6 --- .../executor/_line_execution_process_pool.py | 51 ++++--------------- .../unittests/executor/test_flow_executor.py | 48 ++++------------- 3 files changed, 22 insertions(+), 83 deletions(-) diff --git a/src/promptflow/promptflow/executor/_errors.py b/src/promptflow/promptflow/executor/_errors.py index 321627d3522..9ec5b3d10c2 100644 --- a/src/promptflow/promptflow/executor/_errors.py +++ b/src/promptflow/promptflow/executor/_errors.py @@ -175,12 +175,6 @@ def __init__(self, line_number, timeout): ) -class InvalidMemoryUsageFactor(UserErrorException): - """Exception raised when node reference not found or unsupported""" - - pass - - class EmptyLLMApiMapping(UserErrorException): """Exception raised when connection_type_to_api_mapping is empty and llm node provider can't be inferred""" diff --git a/src/promptflow/promptflow/executor/_line_execution_process_pool.py b/src/promptflow/promptflow/executor/_line_execution_process_pool.py index 3c791a33d9a..a4af9f29503 100644 --- a/src/promptflow/promptflow/executor/_line_execution_process_pool.py +++ b/src/promptflow/promptflow/executor/_line_execution_process_pool.py @@ -27,7 +27,7 @@ from promptflow.contracts.run_info import RunInfo as NodeRunInfo from promptflow.contracts.run_info import Status from promptflow.exceptions import ErrorTarget, PromptflowException -from promptflow.executor._errors import LineExecutionTimeoutError, InvalidMemoryUsageFactor +from promptflow.executor._errors import LineExecutionTimeoutError from promptflow.executor._result import LineResult from promptflow.executor.flow_executor import DEFAULT_CONCURRENCY_BULK, FlowExecutor from promptflow.storage import AbstractRunStorage @@ -197,19 +197,18 @@ def __enter__(self): # Starting a new process in non-fork mode requires to allocate memory. Determine the maximum number of processes # based on available memory to avoid memory bursting. if not self._use_fork: - memory_usage_factor = os.environ.get("PF_BATCH_Memory_Usage_Factor") - available_max_worker_count = get_available_max_worker_count(memory_usage_factor) + available_max_worker_count = get_available_max_worker_count() self._n_process = min(self._worker_count, self._nlines, available_max_worker_count) bulk_logger.info( - f"Not using fork, the number of processes is determined by calculating the difference between the" - f"current system available memory and the total memory usage factor set by the environment" - f"variable multiplied by the total memory, and then dividing this difference by the memory of process." - f"If not set memory usage factor, the default is 1, process count: {self._n_process}") + f"Not using fork, calculate the current system available memory divided by process memory, " + f"and take the minimum value of this value, the worker_count set by the environment variable " + f"configuration, and the row count as the final number of processes. " + f"process count: {self._n_process}") else: self._n_process = min(self._worker_count, self._nlines) bulk_logger.info( - f"Using fork, the number of processes is determined by the lesser of the worker_count set by the" - f"environment variable configuration and the row count, If not set the worker_count, the default" + f"Using fork, the number of processes is determined by the lesser of the worker_count set by the " + f"environment variable configuration and the row count, If not set the worker_count, the default " f"is 16, process count: {self._n_process}") pool = ThreadPool(self._n_process, initializer=set_context, initargs=(contextvars.copy_context(),)) self._pool = pool @@ -545,20 +544,14 @@ def create_executor_legacy(*, flow, connections, loaded_tools, cache_manager, st ) -def get_available_max_worker_count(memory_usage_factor=None): - if memory_usage_factor is None: - memory_usage_factor = 1 - memory_usage_factor = convert_to_decimal(memory_usage_factor) +def get_available_max_worker_count(): pid = os.getpid() mem_info = psutil.virtual_memory() - total_memory = mem_info.total / (1024 * 1024) # in MB - sys_available_memory = mem_info.available / (1024 * 1024) # in MB + available_memory = mem_info.available / (1024 * 1024) # in MB process = psutil.Process(pid) process_memory_info = process.memory_info() process_memory = process_memory_info.rss / (1024 * 1024) # in MB - # To ensure system stability, reserve memory for system usage. - available_memory = (sys_available_memory - (1-memory_usage_factor) * total_memory) - available_max_worker_count = math.floor(available_memory / process_memory) + available_max_worker_count = math.floor((available_memory) / process_memory) if available_max_worker_count < 1: # TODO: For the case of vector db, Optimize execution logic # 1. Let the main process not consume memory because it does not actually invoke @@ -582,25 +575,3 @@ def get_multiprocessing_context(multiprocessing_start_method=None): else: context = multiprocessing.get_context() return context - - -def convert_to_decimal(memory_usage_factor): - if isinstance(memory_usage_factor, (int, float)): - return memory_usage_factor - elif memory_usage_factor.endswith('%'): - # Remove the '%' sign and convert to float - number = float(memory_usage_factor.rstrip('%')) - # Convert the percentage to a decimal - return number / 100 - else: - try: - # Try to convert the string to a float - return float(memory_usage_factor) - except ValueError: - # If the string cannot be converted to a float, raise an error - raise InvalidMemoryUsageFactor( - message_format=( - f"Cannot convert string to number: {memory_usage_factor}" - ), - memory_usage_factor=memory_usage_factor - ) diff --git a/src/promptflow/tests/executor/unittests/executor/test_flow_executor.py b/src/promptflow/tests/executor/unittests/executor/test_flow_executor.py index b0f5e167e12..ba38ae1afac 100644 --- a/src/promptflow/tests/executor/unittests/executor/test_flow_executor.py +++ b/src/promptflow/tests/executor/unittests/executor/test_flow_executor.py @@ -6,7 +6,6 @@ from promptflow.contracts.flow import FlowInputDefinition from promptflow.contracts.tool import ValueType from promptflow.executor._line_execution_process_pool import get_available_max_worker_count -from promptflow.executor._errors import InvalidMemoryUsageFactor from promptflow.executor.flow_executor import ( FlowExecutor, @@ -184,50 +183,25 @@ def test_non_streaming_tool_should_not_be_affected(self): class TestGetAvailableMaxWorkerCount: @pytest.mark.parametrize( - ( - "total_memory", - "available_memory", - "process_memory", - "memory_usage_factor", - "expected_max_worker_count", - "actual_calculate_worker_count", - "expect_exception" - ), + "available_memory, process_memory, expected_max_worker_count, actual_calculate_worker_count", [ - (1024.0, 128.0, 64.0, 0, 1, -14, False), # available_memory - memory_usage_factor * total_memory < 0 - (1024.0, 512.0, 64.0, 0.5, 1, 0, False), # available_memory - memory_usage_factor * total_memory = 0 - (1024.0, 768.0, 64.0, 0.7, 7, 7, False), # available_memory - memory_usage_factor * total_memory > 0 - (1024.0, 768.0, 64.0, 0.8, 8, 8, False), # memory_usage_factor is None - (1024.0, 512.0, 64.0, "50%", 1, 0, False), # The environment variable is configured as a percentage - (1024.0, 512.0, 64.0, "0.5", 1, 0, False), # The environment variable is configured as a string - (1024.0, 512.0, 64.0, "0.5a", 1, 0, True) + (128.0, 64.0, 2, 2), # available_memory > 0 + (0, 64.0, 1, 0), # available_memory = 0 ], ) def test_get_available_max_worker_count( - self, - total_memory, - available_memory, - process_memory, - memory_usage_factor, - expected_max_worker_count, - actual_calculate_worker_count, - expect_exception + self, available_memory, process_memory, expected_max_worker_count, actual_calculate_worker_count ): with patch("psutil.virtual_memory") as mock_mem: - mock_mem.return_value.total = total_memory * 1024 * 1024 mock_mem.return_value.available = available_memory * 1024 * 1024 with patch("psutil.Process") as mock_process: mock_process.return_value.memory_info.return_value.rss = process_memory * 1024 * 1024 with patch("promptflow.executor._line_execution_process_pool.bulk_logger") as mock_logger: mock_logger.warning.return_value = None - if expect_exception: - with pytest.raises(InvalidMemoryUsageFactor): - get_available_max_worker_count(memory_usage_factor) - else: - max_worker_count = get_available_max_worker_count(memory_usage_factor) - assert max_worker_count == expected_max_worker_count - if actual_calculate_worker_count < 1: - mock_logger.warning.assert_called_with( - f"Available max worker count {actual_calculate_worker_count} is less than 1, " - "set it to 1." - ) + max_worker_count = get_available_max_worker_count() + assert max_worker_count == expected_max_worker_count + if actual_calculate_worker_count < 1: + mock_logger.warning.assert_called_with( + f"Available max worker count {actual_calculate_worker_count} is less than 1, " + "set it to 1." + ) From 864443ab282ff9ef34d936ae58cd0c28ae28f7d2 Mon Sep 17 00:00:00 2001 From: LuLu Zuo Date: Fri, 1 Dec 2023 08:41:21 +0800 Subject: [PATCH 04/32] Add os.environ.get("PF_WORKER_COUNT") --- .../executor/_line_execution_process_pool.py | 40 ++++++++++++++----- 1 file changed, 29 insertions(+), 11 deletions(-) diff --git a/src/promptflow/promptflow/executor/_line_execution_process_pool.py b/src/promptflow/promptflow/executor/_line_execution_process_pool.py index a4af9f29503..cb9d00e3327 100644 --- a/src/promptflow/promptflow/executor/_line_execution_process_pool.py +++ b/src/promptflow/promptflow/executor/_line_execution_process_pool.py @@ -151,6 +151,7 @@ def __init__( self._validate_inputs = validate_inputs self._worker_count = flow_executor._worker_count multiprocessing_start_method = os.environ.get("PF_BATCH_METHOD") + self._pf_worker_count = os.environ.get("PF_WORKER_COUNT") sys_start_methods = multiprocessing.get_all_start_methods() if multiprocessing_start_method and multiprocessing_start_method not in sys_start_methods: bulk_logger.warning( @@ -198,18 +199,35 @@ def __enter__(self): # based on available memory to avoid memory bursting. if not self._use_fork: available_max_worker_count = get_available_max_worker_count() - self._n_process = min(self._worker_count, self._nlines, available_max_worker_count) - bulk_logger.info( - f"Not using fork, calculate the current system available memory divided by process memory, " - f"and take the minimum value of this value, the worker_count set by the environment variable " - f"configuration, and the row count as the final number of processes. " - f"process count: {self._n_process}") + if self._pf_worker_count is None: + self._n_process = min(self._worker_count, self._nlines, available_max_worker_count) + bulk_logger.info( + f"Not using fork, and the environment variable PF_WORKER_COUNT is not set. Calculate the current " + f"system available memory divided by process memory, and take the minimum value of this value, " + f"the default value for worker_count 16 and the row count as the final number of processes. " + f"process count: {self._n_process}") + else: + self._n_process = self._pf_worker_count + if available_max_worker_count < self._pf_worker_count: + bulk_logger.warning( + f"The maximum number of processes calculated based on the system available memory " + f"is {available_max_worker_count}, and the PF_WORKER_COUNT is set to {self._pf_worker_count}. " + f"Use the PF_WORKER_COUNT:{self._pf_worker_count} as the final number of processes. " + f"process count: {self._n_process}") + else: + bulk_logger.info( + "PF_WORKER_COUNT:{self._pf_worker_count}, process count: {self._n_process}") else: - self._n_process = min(self._worker_count, self._nlines) - bulk_logger.info( - f"Using fork, the number of processes is determined by the lesser of the worker_count set by the " - f"environment variable configuration and the row count, If not set the worker_count, the default " - f"is 16, process count: {self._n_process}") + if self._pf_worker_count is None: + self._n_process = min(self._worker_count, self._nlines) + bulk_logger.info( + f"Using fork, and the environment variable PF_WORKER_COUNT is not set. The number of processes is " + f"determined by the lesser of the default value for worker_count 16 and the row count." + f"process count: {self._n_process}") + else: + self._n_process = self._pf_worker_count + bulk_logger.info( + "PF_WORKER_COUNT:{self._pf_worker_count}, process count: {self._n_process}") pool = ThreadPool(self._n_process, initializer=set_context, initargs=(contextvars.copy_context(),)) self._pool = pool From fa4196cf94b2a7ceffc3184aa03e632ffb70d573 Mon Sep 17 00:00:00 2001 From: LuLu Zuo Date: Fri, 1 Dec 2023 16:53:02 +0800 Subject: [PATCH 05/32] Add some test case and optimize logic --- .../executor/_line_execution_process_pool.py | 18 +-- .../promptflow/executor/flow_executor.py | 8 +- .../test_line_execution_process_pool.py | 115 ++++++++++++++++++ 3 files changed, 132 insertions(+), 9 deletions(-) diff --git a/src/promptflow/promptflow/executor/_line_execution_process_pool.py b/src/promptflow/promptflow/executor/_line_execution_process_pool.py index cb9d00e3327..1057fe6a264 100644 --- a/src/promptflow/promptflow/executor/_line_execution_process_pool.py +++ b/src/promptflow/promptflow/executor/_line_execution_process_pool.py @@ -151,7 +151,8 @@ def __init__( self._validate_inputs = validate_inputs self._worker_count = flow_executor._worker_count multiprocessing_start_method = os.environ.get("PF_BATCH_METHOD") - self._pf_worker_count = os.environ.get("PF_WORKER_COUNT") + self._DEFAULT_WORKER_COUNT = flow_executor._DEFAULT_WORKER_COUNT + self._pf_worker_count = flow_executor._pf_worker_count sys_start_methods = multiprocessing.get_all_start_methods() if multiprocessing_start_method and multiprocessing_start_method not in sys_start_methods: bulk_logger.warning( @@ -203,9 +204,10 @@ def __enter__(self): self._n_process = min(self._worker_count, self._nlines, available_max_worker_count) bulk_logger.info( f"Not using fork, and the environment variable PF_WORKER_COUNT is not set. Calculate the current " - f"system available memory divided by process memory, and take the minimum value of this value, " - f"the default value for worker_count 16 and the row count as the final number of processes. " - f"process count: {self._n_process}") + f"system available memory divided by process memory, and take the minimum value of this value: " + f"{available_max_worker_count}, the default value for worker_count: {self._DEFAULT_WORKER_COUNT} " + f"and the row count: {self._nlines} as the final number of processes. process count: " + f"{self._n_process}") else: self._n_process = self._pf_worker_count if available_max_worker_count < self._pf_worker_count: @@ -216,18 +218,18 @@ def __enter__(self): f"process count: {self._n_process}") else: bulk_logger.info( - "PF_WORKER_COUNT:{self._pf_worker_count}, process count: {self._n_process}") + f"PF_WORKER_COUNT:{self._pf_worker_count}, process count: {self._n_process}") else: if self._pf_worker_count is None: self._n_process = min(self._worker_count, self._nlines) bulk_logger.info( f"Using fork, and the environment variable PF_WORKER_COUNT is not set. The number of processes is " - f"determined by the lesser of the default value for worker_count 16 and the row count." - f"process count: {self._n_process}") + f"determined by the lesser of the default value for worker_count {self._DEFAULT_WORKER_COUNT} and " + f"the row count: {self._nlines}. process count: {self._n_process}") else: self._n_process = self._pf_worker_count bulk_logger.info( - "PF_WORKER_COUNT:{self._pf_worker_count}, process count: {self._n_process}") + f"PF_WORKER_COUNT:{self._pf_worker_count}, process count: {self._n_process}") pool = ThreadPool(self._n_process, initializer=set_context, initargs=(contextvars.copy_context(),)) self._pool = pool diff --git a/src/promptflow/promptflow/executor/flow_executor.py b/src/promptflow/promptflow/executor/flow_executor.py index 4160058ce16..d6a0e8540c4 100644 --- a/src/promptflow/promptflow/executor/flow_executor.py +++ b/src/promptflow/promptflow/executor/flow_executor.py @@ -130,7 +130,13 @@ def __init__( self._worker_count = worker_count else: try: - worker_count = int(os.environ.get("PF_WORKER_COUNT", self._DEFAULT_WORKER_COUNT)) + pf_worker_count = os.environ.get("PF_WORKER_COUNT") + if pf_worker_count is None: + self._pf_worker_count = pf_worker_count + worker_count = self._DEFAULT_WORKER_COUNT + else: + self._pf_worker_count = int(pf_worker_count) + worker_count = self._pf_worker_count self._worker_count = worker_count except Exception: self._worker_count = self._DEFAULT_WORKER_COUNT diff --git a/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py b/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py index 1fed0c670c5..0455368efbf 100644 --- a/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py +++ b/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py @@ -258,3 +258,118 @@ def test_process_pool_run_with_exception(self, flow_folder, dev_connections, moc assert e.value.message == test_error_msg assert e.value.target == ErrorTarget.AZURE_RUN_STORAGE assert e.value.error_codes[0] == "UserError" + + @pytest.mark.parametrize( + ( + "flow_folder", + "is_set_environ_pf_worker_count", + "pf_worker_count", + "n_process" + ), + [ + (SAMPLE_FLOW, True, "3", 3), + (SAMPLE_FLOW, False, None, 4) + ], + ) + def test_process_pool_parallelism_in_fork_mode( + self, + dev_connections, + flow_folder, + is_set_environ_pf_worker_count, + pf_worker_count, + n_process): + if is_set_environ_pf_worker_count: + os.environ["PF_WORKER_COUNT"] = pf_worker_count + executor = FlowExecutor.create( + get_yaml_file(flow_folder), + dev_connections, + ) + run_id = str(uuid.uuid4()) + bulk_inputs = self.get_bulk_inputs() + nlines = len(bulk_inputs) + + with patch("promptflow.executor._line_execution_process_pool.bulk_logger") as mock_logger: + with LineExecutionProcessPool( + executor, + nlines, + run_id, + "", + False, + None, + ) as pool: + assert pool._n_process == n_process + if is_set_environ_pf_worker_count: + mock_logger.info.assert_called_with( + f"PF_WORKER_COUNT:{pf_worker_count}, process count: {n_process}") + else: + mock_logger.info.assert_called_with( + f"Using fork, and the environment variable PF_WORKER_COUNT is not set. The number of processes " + f"is determined by the lesser of the default value for worker_count " + f"{pool._DEFAULT_WORKER_COUNT} and the row count: {nlines}. process count: {n_process}") + + @pytest.mark.parametrize( + ( + "flow_folder", + "is_set_environ_pf_worker_count", + "is_calculation_smaller_than_set", + "pf_worker_count", + "available_max_worker_count", + "n_process" + ), + [ + (SAMPLE_FLOW, True, False, "2", 4, 2), + (SAMPLE_FLOW, True, True, "6", 2, 6), + (SAMPLE_FLOW, False, True, None, 2, 2) + ], + ) + def test_process_pool_parallelism_in_non_fork_mode( + self, + dev_connections, + flow_folder, + is_set_environ_pf_worker_count, + is_calculation_smaller_than_set, + pf_worker_count, + available_max_worker_count, + n_process + ): + os.environ["PF_BATCH_METHOD"] = "spawn" + if is_set_environ_pf_worker_count: + os.environ["PF_WORKER_COUNT"] = pf_worker_count + executor = FlowExecutor.create( + get_yaml_file(flow_folder), + dev_connections, + ) + run_id = str(uuid.uuid4()) + bulk_inputs = self.get_bulk_inputs() + nlines = len(bulk_inputs) + + with patch("psutil.virtual_memory") as mock_mem: + mock_mem.return_value.available = 128.0 * 1024 * 1024 + with patch("psutil.Process") as mock_process: + mock_process.return_value.memory_info.return_value.rss = 64 * 1024 * 1024 + with patch("promptflow.executor._line_execution_process_pool.bulk_logger") as mock_logger: + with LineExecutionProcessPool( + executor, + nlines, + run_id, + "", + False, + None, + ) as pool: + assert pool._n_process == n_process + if is_set_environ_pf_worker_count and is_calculation_smaller_than_set: + mock_logger.warning.assert_called_with( + f"The maximum number of processes calculated based on the system available memory " + f"is {available_max_worker_count}, and the PF_WORKER_COUNT is set to {pf_worker_count}" + f". Use the PF_WORKER_COUNT:{pf_worker_count} as the final number of processes. " + f"process count: {n_process}") + elif is_set_environ_pf_worker_count and not is_calculation_smaller_than_set: + mock_logger.info.assert_called_with( + f"PF_WORKER_COUNT:{pf_worker_count}, process count: {n_process}") + elif not is_set_environ_pf_worker_count: + mock_logger.info.assert_called_with( + f"Not using fork, and the environment variable PF_WORKER_COUNT is not set. Calculate " + f"the current system available memory divided by process memory, and take the minimum " + f"value of this value: {available_max_worker_count}, the default value for worker_count" + f": {pool._DEFAULT_WORKER_COUNT} and the row count: {nlines} as the final number of " + f"processes. process count: {n_process}") From 8daab14ebceb31e11b2b1498da604e3c1599c99c Mon Sep 17 00:00:00 2001 From: LuLu Zuo Date: Sun, 3 Dec 2023 22:59:57 +0800 Subject: [PATCH 06/32] Delete the previous environment variable --- .../unittests/processpool/test_line_execution_process_pool.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py b/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py index 0455368efbf..c5f31aa426d 100644 --- a/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py +++ b/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py @@ -278,6 +278,7 @@ def test_process_pool_parallelism_in_fork_mode( is_set_environ_pf_worker_count, pf_worker_count, n_process): + os.environ.pop("PF_WORKER_COUNT", None) # Delete the environment variable if it exists if is_set_environ_pf_worker_count: os.environ["PF_WORKER_COUNT"] = pf_worker_count executor = FlowExecutor.create( @@ -332,6 +333,7 @@ def test_process_pool_parallelism_in_non_fork_mode( available_max_worker_count, n_process ): + os.environ.pop("PF_WORKER_COUNT", None) # Delete the environment variable if it exists os.environ["PF_BATCH_METHOD"] = "spawn" if is_set_environ_pf_worker_count: os.environ["PF_WORKER_COUNT"] = pf_worker_count From a8af8fe928e54187e629f19398fa00d31b46b4b8 Mon Sep 17 00:00:00 2001 From: LuLu Zuo Date: Mon, 4 Dec 2023 10:58:13 +0800 Subject: [PATCH 07/32] Delete the previous environment variable --- .../unittests/processpool/test_line_execution_process_pool.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py b/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py index c5f31aa426d..8fd8d972233 100644 --- a/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py +++ b/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py @@ -278,7 +278,6 @@ def test_process_pool_parallelism_in_fork_mode( is_set_environ_pf_worker_count, pf_worker_count, n_process): - os.environ.pop("PF_WORKER_COUNT", None) # Delete the environment variable if it exists if is_set_environ_pf_worker_count: os.environ["PF_WORKER_COUNT"] = pf_worker_count executor = FlowExecutor.create( @@ -307,6 +306,7 @@ def test_process_pool_parallelism_in_fork_mode( f"Using fork, and the environment variable PF_WORKER_COUNT is not set. The number of processes " f"is determined by the lesser of the default value for worker_count " f"{pool._DEFAULT_WORKER_COUNT} and the row count: {nlines}. process count: {n_process}") + os.environ.pop("PF_WORKER_COUNT", None) # Delete the environment variable if it exists @pytest.mark.parametrize( ( @@ -333,7 +333,6 @@ def test_process_pool_parallelism_in_non_fork_mode( available_max_worker_count, n_process ): - os.environ.pop("PF_WORKER_COUNT", None) # Delete the environment variable if it exists os.environ["PF_BATCH_METHOD"] = "spawn" if is_set_environ_pf_worker_count: os.environ["PF_WORKER_COUNT"] = pf_worker_count @@ -375,3 +374,4 @@ def test_process_pool_parallelism_in_non_fork_mode( f"value of this value: {available_max_worker_count}, the default value for worker_count" f": {pool._DEFAULT_WORKER_COUNT} and the row count: {nlines} as the final number of " f"processes. process count: {n_process}") + os.environ.pop("PF_WORKER_COUNT", None) # Delete the environment variable if it exists From 279a8e6bb301e34aae70247e4b7ad892f6d21d69 Mon Sep 17 00:00:00 2001 From: LuLu Zuo Date: Mon, 4 Dec 2023 15:59:13 +0800 Subject: [PATCH 08/32] Optimize log information --- .../executor/_line_execution_process_pool.py | 52 ++-- .../promptflow/executor/flow_executor.py | 6 +- .../unittests/executor/test_flow_executor.py | 27 -- .../test_line_execution_process_pool.py | 270 +++++++++++------- 4 files changed, 206 insertions(+), 149 deletions(-) diff --git a/src/promptflow/promptflow/executor/_line_execution_process_pool.py b/src/promptflow/promptflow/executor/_line_execution_process_pool.py index 1057fe6a264..916ddb8a071 100644 --- a/src/promptflow/promptflow/executor/_line_execution_process_pool.py +++ b/src/promptflow/promptflow/executor/_line_execution_process_pool.py @@ -152,7 +152,7 @@ def __init__( self._worker_count = flow_executor._worker_count multiprocessing_start_method = os.environ.get("PF_BATCH_METHOD") self._DEFAULT_WORKER_COUNT = flow_executor._DEFAULT_WORKER_COUNT - self._pf_worker_count = flow_executor._pf_worker_count + self._use_default_worker_count = flow_executor._use_default_worker_count sys_start_methods = multiprocessing.get_all_start_methods() if multiprocessing_start_method and multiprocessing_start_method not in sys_start_methods: bulk_logger.warning( @@ -199,37 +199,41 @@ def __enter__(self): # Starting a new process in non-fork mode requires to allocate memory. Determine the maximum number of processes # based on available memory to avoid memory bursting. if not self._use_fork: - available_max_worker_count = get_available_max_worker_count() - if self._pf_worker_count is None: - self._n_process = min(self._worker_count, self._nlines, available_max_worker_count) + estimated_available_worker_count = get_available_max_worker_count() + if self._use_default_worker_count: + self._n_process = min(self._worker_count, self._nlines, estimated_available_worker_count) bulk_logger.info( f"Not using fork, and the environment variable PF_WORKER_COUNT is not set. Calculate the current " f"system available memory divided by process memory, and take the minimum value of this value: " - f"{available_max_worker_count}, the default value for worker_count: {self._DEFAULT_WORKER_COUNT} " - f"and the row count: {self._nlines} as the final number of processes. process count: " - f"{self._n_process}") + f"{estimated_available_worker_count}, the default value for worker_count: " + f"{self._DEFAULT_WORKER_COUNT} and the row count: {self._nlines} as the final number of processes. " + f"process count: {self._n_process}") else: - self._n_process = self._pf_worker_count - if available_max_worker_count < self._pf_worker_count: + self._n_process = self._worker_count + if estimated_available_worker_count < self._worker_count: bulk_logger.warning( f"The maximum number of processes calculated based on the system available memory " - f"is {available_max_worker_count}, and the PF_WORKER_COUNT is set to {self._pf_worker_count}. " - f"Use the PF_WORKER_COUNT:{self._pf_worker_count} as the final number of processes. " + f"is {estimated_available_worker_count}, and the PF_WORKER_COUNT is set to {self._worker_count}. " + f"Use the PF_WORKER_COUNT:{self._worker_count} as the final number of processes. " f"process count: {self._n_process}") else: bulk_logger.info( - f"PF_WORKER_COUNT:{self._pf_worker_count}, process count: {self._n_process}") + f"The number of processes has been explicitly set to {self._worker_count} by the " + f"environment variable PF_WORKER_COUNT. This value will determine the degree of parallelism of " + f"the process. process count: {self._n_process}") else: - if self._pf_worker_count is None: + if self._use_default_worker_count: self._n_process = min(self._worker_count, self._nlines) bulk_logger.info( f"Using fork, and the environment variable PF_WORKER_COUNT is not set. The number of processes is " f"determined by the lesser of the default value for worker_count {self._DEFAULT_WORKER_COUNT} and " f"the row count: {self._nlines}. process count: {self._n_process}") else: - self._n_process = self._pf_worker_count + self._n_process = self._worker_count bulk_logger.info( - f"PF_WORKER_COUNT:{self._pf_worker_count}, process count: {self._n_process}") + f"The number of processes has been explicitly set to {self._worker_count} by the " + f"environment variable PF_WORKER_COUNT. This value will determine the degree of parallelism of " + f"the process. process count: {self._n_process}") pool = ThreadPool(self._n_process, initializer=set_context, initargs=(contextvars.copy_context(),)) self._pool = pool @@ -571,20 +575,22 @@ def get_available_max_worker_count(): process = psutil.Process(pid) process_memory_info = process.memory_info() process_memory = process_memory_info.rss / (1024 * 1024) # in MB - available_max_worker_count = math.floor((available_memory) / process_memory) - if available_max_worker_count < 1: + estimated_available_worker_count = math.floor((available_memory) / process_memory) + if estimated_available_worker_count < 1: # TODO: For the case of vector db, Optimize execution logic # 1. Let the main process not consume memory because it does not actually invoke # 2. When the degree of parallelism is 1, main process executes the task directly and not # create the child process - bulk_logger.warning(f"Available max worker count {available_max_worker_count} is less than 1, set it to 1.") - available_max_worker_count = 1 + bulk_logger.warning( + f"Available max worker count {estimated_available_worker_count} is less than 1, set it to 1.") + estimated_available_worker_count = 1 bulk_logger.info( - f"""Process {pid} current available memory is {process_memory}, - memory consumption of current process is {available_memory}, - worker count is set to {available_memory}/{process_memory} = {available_max_worker_count}""" + f"Process {pid} current available memory is {process_memory}, " + f"memory consumption of current process is {available_memory}, " + f"the estimated available worker count is {available_memory}/{process_memory} " + f"= {estimated_available_worker_count}" ) - return available_max_worker_count + return estimated_available_worker_count def get_multiprocessing_context(multiprocessing_start_method=None): diff --git a/src/promptflow/promptflow/executor/flow_executor.py b/src/promptflow/promptflow/executor/flow_executor.py index d6a0e8540c4..02164eb456f 100644 --- a/src/promptflow/promptflow/executor/flow_executor.py +++ b/src/promptflow/promptflow/executor/flow_executor.py @@ -132,11 +132,11 @@ def __init__( try: pf_worker_count = os.environ.get("PF_WORKER_COUNT") if pf_worker_count is None: - self._pf_worker_count = pf_worker_count + self._use_default_worker_count = True worker_count = self._DEFAULT_WORKER_COUNT else: - self._pf_worker_count = int(pf_worker_count) - worker_count = self._pf_worker_count + self._use_default_worker_count = False + worker_count = int(pf_worker_count) self._worker_count = worker_count except Exception: self._worker_count = self._DEFAULT_WORKER_COUNT diff --git a/src/promptflow/tests/executor/unittests/executor/test_flow_executor.py b/src/promptflow/tests/executor/unittests/executor/test_flow_executor.py index ba38ae1afac..f47dc40a5cd 100644 --- a/src/promptflow/tests/executor/unittests/executor/test_flow_executor.py +++ b/src/promptflow/tests/executor/unittests/executor/test_flow_executor.py @@ -5,7 +5,6 @@ from promptflow import tool from promptflow.contracts.flow import FlowInputDefinition from promptflow.contracts.tool import ValueType -from promptflow.executor._line_execution_process_pool import get_available_max_worker_count from promptflow.executor.flow_executor import ( FlowExecutor, @@ -179,29 +178,3 @@ def test_streaming_tool_should_be_consumed_and_merged(self): def test_non_streaming_tool_should_not_be_affected(self): func = _ensure_node_result_is_serializable(non_streaming_tool) assert func() == 1 - - -class TestGetAvailableMaxWorkerCount: - @pytest.mark.parametrize( - "available_memory, process_memory, expected_max_worker_count, actual_calculate_worker_count", - [ - (128.0, 64.0, 2, 2), # available_memory > 0 - (0, 64.0, 1, 0), # available_memory = 0 - ], - ) - def test_get_available_max_worker_count( - self, available_memory, process_memory, expected_max_worker_count, actual_calculate_worker_count - ): - with patch("psutil.virtual_memory") as mock_mem: - mock_mem.return_value.available = available_memory * 1024 * 1024 - with patch("psutil.Process") as mock_process: - mock_process.return_value.memory_info.return_value.rss = process_memory * 1024 * 1024 - with patch("promptflow.executor._line_execution_process_pool.bulk_logger") as mock_logger: - mock_logger.warning.return_value = None - max_worker_count = get_available_max_worker_count() - assert max_worker_count == expected_max_worker_count - if actual_calculate_worker_count < 1: - mock_logger.warning.assert_called_with( - f"Available max worker count {actual_calculate_worker_count} is less than 1, " - "set it to 1." - ) diff --git a/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py b/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py index 8fd8d972233..8251b789574 100644 --- a/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py +++ b/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py @@ -17,6 +17,7 @@ LineExecutionProcessPool, _exec_line, get_multiprocessing_context, + get_available_max_worker_count ) from promptflow.executor._result import LineResult @@ -25,31 +26,129 @@ SAMPLE_FLOW = "web_classification_no_variants" -@pytest.mark.unittest -class TestLineExecutionProcessPool: - def get_line_inputs(self, flow_folder=""): - if flow_folder: - inputs = self.get_bulk_inputs(flow_folder) - return inputs[0] - return { - "url": "https://www.microsoft.com/en-us/windows/", - "text": "some_text", - } +def get_line_inputs(flow_folder=""): + if flow_folder: + inputs = get_bulk_inputs(flow_folder) + return inputs[0] + return { + "url": "https://www.microsoft.com/en-us/windows/", + "text": "some_text", + } + - def get_bulk_inputs(self, nlinee=4, flow_folder="", sample_inputs_file="", return_dict=False): - if flow_folder: - if not sample_inputs_file: - sample_inputs_file = "samples.json" - inputs = get_flow_sample_inputs(flow_folder, sample_inputs_file=sample_inputs_file) - if isinstance(inputs, list) and len(inputs) > 0: +def get_bulk_inputs(nlinee=4, flow_folder="", sample_inputs_file="", return_dict=False): + if flow_folder: + if not sample_inputs_file: + sample_inputs_file = "samples.json" + inputs = get_flow_sample_inputs(flow_folder, sample_inputs_file=sample_inputs_file) + if isinstance(inputs, list) and len(inputs) > 0: + return inputs + elif isinstance(inputs, dict): + if return_dict: return inputs - elif isinstance(inputs, dict): - if return_dict: - return inputs - return [inputs] + return [inputs] + else: + raise Exception(f"Invalid type of bulk input: {inputs}") + return [get_line_inputs() for _ in range(nlinee)] + + +@pytest.mark.skip("This is a subprocess function used for testing and cannot be tested alone.") +def test_fork_mode_parallelism_in_subprocess( + dev_connections, + flow_folder, + is_set_environ_pf_worker_count, + pf_worker_count, + n_process): + + if is_set_environ_pf_worker_count: + os.environ["PF_WORKER_COUNT"] = pf_worker_count + executor = FlowExecutor.create( + get_yaml_file(flow_folder), + dev_connections, + ) + run_id = str(uuid.uuid4()) + bulk_inputs = get_bulk_inputs() + nlines = len(bulk_inputs) + + with patch("promptflow.executor._line_execution_process_pool.bulk_logger") as mock_logger: + with LineExecutionProcessPool( + executor, + nlines, + run_id, + "", + False, + None, + ) as pool: + assert pool._n_process == n_process + if is_set_environ_pf_worker_count: + mock_logger.info.assert_called_with( + f"The number of processes has been explicitly set to {pf_worker_count} by the " + f"environment variable PF_WORKER_COUNT. This value will determine the degree of " + f"parallelism of the process. process count: {n_process}") else: - raise Exception(f"Invalid type of bulk input: {inputs}") - return [self.get_line_inputs() for _ in range(nlinee)] + mock_logger.info.assert_called_with( + f"Using fork, and the environment variable PF_WORKER_COUNT is not set. The number of processes " + f"is determined by the lesser of the default value for worker_count " + f"{pool._DEFAULT_WORKER_COUNT} and the row count: {nlines}. process count: {n_process}") + + +@pytest.mark.skip("This is a subprocess function used for testing and cannot be tested alone.") +def test_spawn_mode_parallelism_in_subprocess( + dev_connections, + flow_folder, + is_set_environ_pf_worker_count, + is_calculation_smaller_than_set, + pf_worker_count, + available_max_worker_count, + n_process +): + os.environ["PF_BATCH_METHOD"] = "spawn" + if is_set_environ_pf_worker_count: + os.environ["PF_WORKER_COUNT"] = pf_worker_count + executor = FlowExecutor.create( + get_yaml_file(flow_folder), + dev_connections, + ) + run_id = str(uuid.uuid4()) + bulk_inputs = get_bulk_inputs() + nlines = len(bulk_inputs) + + with patch("psutil.virtual_memory") as mock_mem: + mock_mem.return_value.available = 128.0 * 1024 * 1024 + with patch("psutil.Process") as mock_process: + mock_process.return_value.memory_info.return_value.rss = 64 * 1024 * 1024 + with patch("promptflow.executor._line_execution_process_pool.bulk_logger") as mock_logger: + with LineExecutionProcessPool( + executor, + nlines, + run_id, + "", + False, + None, + ) as pool: + assert pool._n_process == n_process + if is_set_environ_pf_worker_count and is_calculation_smaller_than_set: + mock_logger.warning.assert_called_with( + f"The maximum number of processes calculated based on the system available memory " + f"is {available_max_worker_count}, and the PF_WORKER_COUNT is set to {pf_worker_count}" + f". Use the PF_WORKER_COUNT:{pf_worker_count} as the final number of processes. " + f"process count: {n_process}") + elif is_set_environ_pf_worker_count and not is_calculation_smaller_than_set: + mock_logger.info.assert_called_with( + f"The number of processes has been explicitly set to {pf_worker_count} by the " + f"environment variable PF_WORKER_COUNT. This value will determine the degree of " + f"parallelism of the process. process count: {n_process}") + elif not is_set_environ_pf_worker_count: + mock_logger.info.assert_called_with( + f"Not using fork, and the environment variable PF_WORKER_COUNT is not set. Calculate " + f"the current system available memory divided by process memory, and take the minimum " + f"value of this value: {available_max_worker_count}, the default value for worker_count" + f": {pool._DEFAULT_WORKER_COUNT} and the row count: {nlines} as the final number of " + f"processes. process count: {n_process}") + + +@pytest.mark.unittest +class TestLineExecutionProcessPool: def create_line_execution_process_pool(self, dev_connections): executor = FlowExecutor.create( @@ -58,7 +157,7 @@ def create_line_execution_process_pool(self, dev_connections): line_timeout_sec=1, ) run_id = str(uuid.uuid4()) - bulk_inputs = self.get_bulk_inputs() + bulk_inputs = get_bulk_inputs() nlines = len(bulk_inputs) line_execution_process_pool = LineExecutionProcessPool( executor, @@ -84,7 +183,7 @@ def test_line_execution_process_pool(self, flow_folder, dev_connections): executor = FlowExecutor.create(get_yaml_file(flow_folder), dev_connections) executor._log_interval = 1 run_id = str(uuid.uuid4()) - bulk_inputs = self.get_bulk_inputs() + bulk_inputs = get_bulk_inputs() nlines = len(bulk_inputs) run_id = run_id or str(uuid.uuid4()) with LineExecutionProcessPool( @@ -114,7 +213,7 @@ def test_line_execution_not_completed(self, flow_folder, dev_connections): line_timeout_sec=1, ) run_id = str(uuid.uuid4()) - bulk_inputs = self.get_bulk_inputs() + bulk_inputs = get_bulk_inputs() nlines = len(bulk_inputs) with LineExecutionProcessPool( executor, @@ -147,7 +246,7 @@ def test_exec_line(self, flow_folder, dev_connections, mocker: MockFixture): line_timeout_sec=1, ) run_id = str(uuid.uuid4()) - line_inputs = self.get_line_inputs() + line_inputs = get_line_inputs() line_result = _exec_line( executor=executor, output_queue=output_queue, @@ -178,7 +277,7 @@ def test_exec_line_failed_when_line_execution_not_start(self, flow_folder, dev_c message=test_error_msg, target=ErrorTarget.AZURE_RUN_STORAGE ) run_id = str(uuid.uuid4()) - line_inputs = self.get_line_inputs() + line_inputs = get_line_inputs() line_result = _exec_line( executor=executor, output_queue=output_queue, @@ -243,7 +342,7 @@ def test_process_pool_run_with_exception(self, flow_folder, dev_connections, moc dev_connections, ) run_id = str(uuid.uuid4()) - bulk_inputs = self.get_bulk_inputs() + bulk_inputs = get_bulk_inputs() nlines = len(bulk_inputs) with LineExecutionProcessPool( executor, @@ -278,35 +377,14 @@ def test_process_pool_parallelism_in_fork_mode( is_set_environ_pf_worker_count, pf_worker_count, n_process): - if is_set_environ_pf_worker_count: - os.environ["PF_WORKER_COUNT"] = pf_worker_count - executor = FlowExecutor.create( - get_yaml_file(flow_folder), - dev_connections, - ) - run_id = str(uuid.uuid4()) - bulk_inputs = self.get_bulk_inputs() - nlines = len(bulk_inputs) - - with patch("promptflow.executor._line_execution_process_pool.bulk_logger") as mock_logger: - with LineExecutionProcessPool( - executor, - nlines, - run_id, - "", - False, - None, - ) as pool: - assert pool._n_process == n_process - if is_set_environ_pf_worker_count: - mock_logger.info.assert_called_with( - f"PF_WORKER_COUNT:{pf_worker_count}, process count: {n_process}") - else: - mock_logger.info.assert_called_with( - f"Using fork, and the environment variable PF_WORKER_COUNT is not set. The number of processes " - f"is determined by the lesser of the default value for worker_count " - f"{pool._DEFAULT_WORKER_COUNT} and the row count: {nlines}. process count: {n_process}") - os.environ.pop("PF_WORKER_COUNT", None) # Delete the environment variable if it exists + p = multiprocessing.Process(target=test_fork_mode_parallelism_in_subprocess, + args=(dev_connections, + flow_folder, + is_set_environ_pf_worker_count, + pf_worker_count, + n_process)) + p.start() + p.join() @pytest.mark.parametrize( ( @@ -323,7 +401,7 @@ def test_process_pool_parallelism_in_fork_mode( (SAMPLE_FLOW, False, True, None, 2, 2) ], ) - def test_process_pool_parallelism_in_non_fork_mode( + def test_process_pool_parallelism_in_non_spawn_mode( self, dev_connections, flow_folder, @@ -333,45 +411,45 @@ def test_process_pool_parallelism_in_non_fork_mode( available_max_worker_count, n_process ): - os.environ["PF_BATCH_METHOD"] = "spawn" - if is_set_environ_pf_worker_count: - os.environ["PF_WORKER_COUNT"] = pf_worker_count - executor = FlowExecutor.create( - get_yaml_file(flow_folder), - dev_connections, - ) - run_id = str(uuid.uuid4()) - bulk_inputs = self.get_bulk_inputs() - nlines = len(bulk_inputs) + p = multiprocessing.Process(target=test_spawn_mode_parallelism_in_subprocess, + args=(dev_connections, + flow_folder, + is_set_environ_pf_worker_count, + is_calculation_smaller_than_set, + pf_worker_count, + available_max_worker_count, + n_process)) + p.start() + p.join() + +class TestGetAvailableMaxWorkerCount: + @pytest.mark.parametrize( + "available_memory, process_memory, expected_max_worker_count, actual_calculate_worker_count", + [ + (128.0, 64.0, 2, 2), # available_memory > 0 + (0.0, 64.0, 1, 0), # available_memory = 0 + ], + ) + def test_get_available_max_worker_count( + self, available_memory, process_memory, expected_max_worker_count, actual_calculate_worker_count + ): with patch("psutil.virtual_memory") as mock_mem: - mock_mem.return_value.available = 128.0 * 1024 * 1024 + mock_mem.return_value.available = available_memory * 1024 * 1024 with patch("psutil.Process") as mock_process: - mock_process.return_value.memory_info.return_value.rss = 64 * 1024 * 1024 + mock_process.return_value.memory_info.return_value.rss = process_memory * 1024 * 1024 with patch("promptflow.executor._line_execution_process_pool.bulk_logger") as mock_logger: - with LineExecutionProcessPool( - executor, - nlines, - run_id, - "", - False, - None, - ) as pool: - assert pool._n_process == n_process - if is_set_environ_pf_worker_count and is_calculation_smaller_than_set: - mock_logger.warning.assert_called_with( - f"The maximum number of processes calculated based on the system available memory " - f"is {available_max_worker_count}, and the PF_WORKER_COUNT is set to {pf_worker_count}" - f". Use the PF_WORKER_COUNT:{pf_worker_count} as the final number of processes. " - f"process count: {n_process}") - elif is_set_environ_pf_worker_count and not is_calculation_smaller_than_set: - mock_logger.info.assert_called_with( - f"PF_WORKER_COUNT:{pf_worker_count}, process count: {n_process}") - elif not is_set_environ_pf_worker_count: - mock_logger.info.assert_called_with( - f"Not using fork, and the environment variable PF_WORKER_COUNT is not set. Calculate " - f"the current system available memory divided by process memory, and take the minimum " - f"value of this value: {available_max_worker_count}, the default value for worker_count" - f": {pool._DEFAULT_WORKER_COUNT} and the row count: {nlines} as the final number of " - f"processes. process count: {n_process}") - os.environ.pop("PF_WORKER_COUNT", None) # Delete the environment variable if it exists + mock_logger.warning.return_value = None + estimated_available_worker_count = get_available_max_worker_count() + assert estimated_available_worker_count == expected_max_worker_count + if actual_calculate_worker_count < 1: + mock_logger.warning.assert_called_with( + f"Available max worker count {actual_calculate_worker_count} is less than 1, " + "set it to 1." + ) + mock_logger.info.assert_called_with( + f"Process {os.getpid()} current available memory is {process_memory}, " + f"memory consumption of current process is {available_memory}, " + f"the estimated available worker count is {available_memory}/{process_memory} " + f"= {estimated_available_worker_count}" + ) From 89a67ed138854f3bbc2fc607071e8dfcae07c197 Mon Sep 17 00:00:00 2001 From: LuLu Zuo Date: Mon, 4 Dec 2023 16:01:06 +0800 Subject: [PATCH 09/32] fix flask8 --- .../promptflow/executor/_line_execution_process_pool.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/promptflow/promptflow/executor/_line_execution_process_pool.py b/src/promptflow/promptflow/executor/_line_execution_process_pool.py index 916ddb8a071..90cad8032e4 100644 --- a/src/promptflow/promptflow/executor/_line_execution_process_pool.py +++ b/src/promptflow/promptflow/executor/_line_execution_process_pool.py @@ -213,9 +213,9 @@ def __enter__(self): if estimated_available_worker_count < self._worker_count: bulk_logger.warning( f"The maximum number of processes calculated based on the system available memory " - f"is {estimated_available_worker_count}, and the PF_WORKER_COUNT is set to {self._worker_count}. " - f"Use the PF_WORKER_COUNT:{self._worker_count} as the final number of processes. " - f"process count: {self._n_process}") + f"is {estimated_available_worker_count}, and the PF_WORKER_COUNT is set to " + f"{self._worker_count}. Use the PF_WORKER_COUNT:{self._worker_count} as the final " + f"number of processes. process count: {self._n_process}") else: bulk_logger.info( f"The number of processes has been explicitly set to {self._worker_count} by the " From 4ac4cea88a0a702fa2ef963e83bfa0bc1fd9a5d3 Mon Sep 17 00:00:00 2001 From: LuLu Zuo Date: Mon, 4 Dec 2023 16:02:34 +0800 Subject: [PATCH 10/32] fix flask8 --- .../tests/executor/unittests/executor/test_flow_executor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/promptflow/tests/executor/unittests/executor/test_flow_executor.py b/src/promptflow/tests/executor/unittests/executor/test_flow_executor.py index f47dc40a5cd..784ca44e6ef 100644 --- a/src/promptflow/tests/executor/unittests/executor/test_flow_executor.py +++ b/src/promptflow/tests/executor/unittests/executor/test_flow_executor.py @@ -1,4 +1,4 @@ -from unittest.mock import Mock, patch +from unittest.mock import Mock import pytest From 9b3d695231942727614c2cfa9f2ea1b7dd4fc918 Mon Sep 17 00:00:00 2001 From: LuLu Zuo Date: Mon, 4 Dec 2023 16:12:11 +0800 Subject: [PATCH 11/32] delete "the" --- .../promptflow/executor/_line_execution_process_pool.py | 2 +- .../unittests/processpool/test_line_execution_process_pool.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/promptflow/promptflow/executor/_line_execution_process_pool.py b/src/promptflow/promptflow/executor/_line_execution_process_pool.py index 90cad8032e4..e002610ffcc 100644 --- a/src/promptflow/promptflow/executor/_line_execution_process_pool.py +++ b/src/promptflow/promptflow/executor/_line_execution_process_pool.py @@ -587,7 +587,7 @@ def get_available_max_worker_count(): bulk_logger.info( f"Process {pid} current available memory is {process_memory}, " f"memory consumption of current process is {available_memory}, " - f"the estimated available worker count is {available_memory}/{process_memory} " + f"estimated available worker count is {available_memory}/{process_memory} " f"= {estimated_available_worker_count}" ) return estimated_available_worker_count diff --git a/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py b/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py index 8251b789574..5108b956772 100644 --- a/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py +++ b/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py @@ -450,6 +450,6 @@ def test_get_available_max_worker_count( mock_logger.info.assert_called_with( f"Process {os.getpid()} current available memory is {process_memory}, " f"memory consumption of current process is {available_memory}, " - f"the estimated available worker count is {available_memory}/{process_memory} " + f"estimated available worker count is {available_memory}/{process_memory} " f"= {estimated_available_worker_count}" ) From 3c61469135b733cd289e60efa05fa6edf1782654 Mon Sep 17 00:00:00 2001 From: LuLu Zuo Date: Mon, 4 Dec 2023 17:40:37 +0800 Subject: [PATCH 12/32] Extract as public function and refine worker_count --- .../executor/_line_execution_process_pool.py | 17 ++++++----- .../promptflow/executor/flow_executor.py | 30 ++++++++++++------- 2 files changed, 28 insertions(+), 19 deletions(-) diff --git a/src/promptflow/promptflow/executor/_line_execution_process_pool.py b/src/promptflow/promptflow/executor/_line_execution_process_pool.py index e002610ffcc..be67dbd9d33 100644 --- a/src/promptflow/promptflow/executor/_line_execution_process_pool.py +++ b/src/promptflow/promptflow/executor/_line_execution_process_pool.py @@ -217,10 +217,7 @@ def __enter__(self): f"{self._worker_count}. Use the PF_WORKER_COUNT:{self._worker_count} as the final " f"number of processes. process count: {self._n_process}") else: - bulk_logger.info( - f"The number of processes has been explicitly set to {self._worker_count} by the " - f"environment variable PF_WORKER_COUNT. This value will determine the degree of parallelism of " - f"the process. process count: {self._n_process}") + log_process_count_info(self._worker_count, self._n_process) else: if self._use_default_worker_count: self._n_process = min(self._worker_count, self._nlines) @@ -230,10 +227,7 @@ def __enter__(self): f"the row count: {self._nlines}. process count: {self._n_process}") else: self._n_process = self._worker_count - bulk_logger.info( - f"The number of processes has been explicitly set to {self._worker_count} by the " - f"environment variable PF_WORKER_COUNT. This value will determine the degree of parallelism of " - f"the process. process count: {self._n_process}") + log_process_count_info(self._worker_count, self._n_process) pool = ThreadPool(self._n_process, initializer=set_context, initargs=(contextvars.copy_context(),)) self._pool = pool @@ -601,3 +595,10 @@ def get_multiprocessing_context(multiprocessing_start_method=None): else: context = multiprocessing.get_context() return context + + +def log_process_count_info(worker_count, process_count): + bulk_logger.info( + f"The number of processes has been explicitly set to {worker_count} by the " + f"environment variable PF_WORKER_COUNT. This value will determine the degree of parallelism of " + f"the process. process count: {process_count}") diff --git a/src/promptflow/promptflow/executor/flow_executor.py b/src/promptflow/promptflow/executor/flow_executor.py index 02164eb456f..6e353a39d76 100644 --- a/src/promptflow/promptflow/executor/flow_executor.py +++ b/src/promptflow/promptflow/executor/flow_executor.py @@ -129,17 +129,14 @@ def __init__( if worker_count is not None: self._worker_count = worker_count else: - try: - pf_worker_count = os.environ.get("PF_WORKER_COUNT") - if pf_worker_count is None: - self._use_default_worker_count = True - worker_count = self._DEFAULT_WORKER_COUNT - else: - self._use_default_worker_count = False - worker_count = int(pf_worker_count) - self._worker_count = worker_count - except Exception: - self._worker_count = self._DEFAULT_WORKER_COUNT + worker_count_in_env = load_worker_count_in_env() + if worker_count_in_env is None: + self._use_default_worker_count = True + worker_count = self._DEFAULT_WORKER_COUNT + else: + self._use_default_worker_count = False + worker_count = int(worker_count_in_env) + self._worker_count = worker_count if self._worker_count <= 0: self._worker_count = self._DEFAULT_WORKER_COUNT self._run_tracker = run_tracker @@ -969,6 +966,17 @@ def ensure_flow_is_serializable(self): self._tools_manager.wrap_tool(node.name, wrapper=_ensure_node_result_is_serializable) +def load_worker_count_in_env(): + try: + pf_worker_count = os.environ.get("PF_WORKER_COUNT") + if pf_worker_count is None: + return None + else: + return int(pf_worker_count) + except Exception: + return None + + def _inject_stream_options(should_stream: Callable[[], bool], streaming_option_parameter="stream"): """Inject the stream options to the decorated function. From 467f2e04cb7444a4a09e841f101d680ec9595fd5 Mon Sep 17 00:00:00 2001 From: LuLu Zuo Date: Mon, 4 Dec 2023 18:25:06 +0800 Subject: [PATCH 13/32] refine load_worker_count_in_env --- .../promptflow/executor/flow_executor.py | 27 ++++++++++--------- 1 file changed, 14 insertions(+), 13 deletions(-) diff --git a/src/promptflow/promptflow/executor/flow_executor.py b/src/promptflow/promptflow/executor/flow_executor.py index 6e353a39d76..b338ef3a980 100644 --- a/src/promptflow/promptflow/executor/flow_executor.py +++ b/src/promptflow/promptflow/executor/flow_executor.py @@ -129,15 +129,10 @@ def __init__( if worker_count is not None: self._worker_count = worker_count else: - worker_count_in_env = load_worker_count_in_env() - if worker_count_in_env is None: - self._use_default_worker_count = True - worker_count = self._DEFAULT_WORKER_COUNT - else: - self._use_default_worker_count = False - worker_count = int(worker_count_in_env) - self._worker_count = worker_count + self.use_default_worker_count, self._worker_count = load_worker_count_in_env(self._DEFAULT_WORKER_COUNT) if self._worker_count <= 0: + logger.warning( + f"Invalid worker count: {self._worker_count}. Resetting to default value: {self._DEFAULT_WORKER_COUNT}") self._worker_count = self._DEFAULT_WORKER_COUNT self._run_tracker = run_tracker self._cache_manager = cache_manager @@ -966,15 +961,21 @@ def ensure_flow_is_serializable(self): self._tools_manager.wrap_tool(node.name, wrapper=_ensure_node_result_is_serializable) -def load_worker_count_in_env(): +def load_worker_count_in_env(DEFAULT_WORKER_COUNT): try: pf_worker_count = os.environ.get("PF_WORKER_COUNT") if pf_worker_count is None: - return None + use_default_worker_count = True + worker_count = DEFAULT_WORKER_COUNT else: - return int(pf_worker_count) - except Exception: - return None + use_default_worker_count = False + worker_count = int(pf_worker_count) + except Exception as e: + logger.warning(f"Failed to convert PF_WORKER_COUNT '{pf_worker_count}' to an integer: {e}") + use_default_worker_count = True + worker_count = DEFAULT_WORKER_COUNT + + return use_default_worker_count, worker_count def _inject_stream_options(should_stream: Callable[[], bool], streaming_option_parameter="stream"): From 1b5a1c3d4ed938d6bbddc384fdad3a74dc18be92 Mon Sep 17 00:00:00 2001 From: LuLu Zuo Date: Mon, 4 Dec 2023 21:49:49 +0800 Subject: [PATCH 14/32] Rename variable --- src/promptflow/promptflow/executor/flow_executor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/promptflow/promptflow/executor/flow_executor.py b/src/promptflow/promptflow/executor/flow_executor.py index b338ef3a980..1c607775553 100644 --- a/src/promptflow/promptflow/executor/flow_executor.py +++ b/src/promptflow/promptflow/executor/flow_executor.py @@ -129,7 +129,7 @@ def __init__( if worker_count is not None: self._worker_count = worker_count else: - self.use_default_worker_count, self._worker_count = load_worker_count_in_env(self._DEFAULT_WORKER_COUNT) + self._use_default_worker_count, self._worker_count = load_worker_count_in_env(self._DEFAULT_WORKER_COUNT) if self._worker_count <= 0: logger.warning( f"Invalid worker count: {self._worker_count}. Resetting to default value: {self._DEFAULT_WORKER_COUNT}") From 8c77dba8824e83de35ac7bc3a435fd4df9349295 Mon Sep 17 00:00:00 2001 From: LuLu Zuo Date: Tue, 5 Dec 2023 11:07:11 +0800 Subject: [PATCH 15/32] Refine load_worker_count_in_env --- .../promptflow/executor/flow_executor.py | 10 +++--- .../unittests/executor/test_flow_executor.py | 33 ++++++++++++++++++- 2 files changed, 38 insertions(+), 5 deletions(-) diff --git a/src/promptflow/promptflow/executor/flow_executor.py b/src/promptflow/promptflow/executor/flow_executor.py index 1c607775553..e4b8ecada88 100644 --- a/src/promptflow/promptflow/executor/flow_executor.py +++ b/src/promptflow/promptflow/executor/flow_executor.py @@ -130,10 +130,6 @@ def __init__( self._worker_count = worker_count else: self._use_default_worker_count, self._worker_count = load_worker_count_in_env(self._DEFAULT_WORKER_COUNT) - if self._worker_count <= 0: - logger.warning( - f"Invalid worker count: {self._worker_count}. Resetting to default value: {self._DEFAULT_WORKER_COUNT}") - self._worker_count = self._DEFAULT_WORKER_COUNT self._run_tracker = run_tracker self._cache_manager = cache_manager self._loaded_tools = loaded_tools @@ -975,6 +971,12 @@ def load_worker_count_in_env(DEFAULT_WORKER_COUNT): use_default_worker_count = True worker_count = DEFAULT_WORKER_COUNT + if worker_count <= 0: + logger.warning( + f"Invalid worker count: {worker_count}. Resetting to default value: {DEFAULT_WORKER_COUNT}") + use_default_worker_count = True + worker_count = DEFAULT_WORKER_COUNT + return use_default_worker_count, worker_count diff --git a/src/promptflow/tests/executor/unittests/executor/test_flow_executor.py b/src/promptflow/tests/executor/unittests/executor/test_flow_executor.py index 784ca44e6ef..a2a5988c55d 100644 --- a/src/promptflow/tests/executor/unittests/executor/test_flow_executor.py +++ b/src/promptflow/tests/executor/unittests/executor/test_flow_executor.py @@ -1,6 +1,7 @@ -from unittest.mock import Mock +from unittest.mock import Mock, patch import pytest +import os from promptflow import tool from promptflow.contracts.flow import FlowInputDefinition @@ -11,7 +12,9 @@ _ensure_node_result_is_serializable, _inject_stream_options, enable_streaming_for_llm_tool, + load_worker_count_in_env ) +from promptflow._utils.logger_utils import logger from promptflow.tools.aoai import chat, completion from promptflow.tools.embedding import embedding @@ -178,3 +181,31 @@ def test_streaming_tool_should_be_consumed_and_merged(self): def test_non_streaming_tool_should_not_be_affected(self): func = _ensure_node_result_is_serializable(non_streaming_tool) assert func() == 1 + + +class TestLoadWorkerCountInEnv: + def test_no_env_var_set(self): + with patch.dict(os.environ, {}, clear=True): + use_default, worker_count = load_worker_count_in_env(4) + assert use_default is True + assert worker_count == 4 + + def test_valid_integer_env_var(self): + with patch.dict(os.environ, {'PF_WORKER_COUNT': '5'}): + use_default, worker_count = load_worker_count_in_env(4) + assert use_default is False + assert worker_count == 5 + + def test_invalid_non_integer_env_var(self): + with patch.dict(os.environ, {'PF_WORKER_COUNT': 'invalid'}), patch.object(logger, 'warning') as mock_warning: + use_default, worker_count = load_worker_count_in_env(4) + mock_warning.assert_called() + assert use_default is True + assert worker_count == 4 + + def test_invalid_non_positive_integer_env_var(self): + with patch.dict(os.environ, {'PF_WORKER_COUNT': '0'}), patch.object(logger, 'warning') as mock_warning: + use_default, worker_count = load_worker_count_in_env(4) + mock_warning.assert_called() + assert use_default is True + assert worker_count == 4 From 8ac1eda6c622631e389249dd8e1f52281f75ac29 Mon Sep 17 00:00:00 2001 From: LuLu Zuo Date: Tue, 5 Dec 2023 12:03:11 +0800 Subject: [PATCH 16/32] Optimize log --- .../executor/_line_execution_process_pool.py | 35 +++++++++---------- .../test_line_execution_process_pool.py | 32 ++++++++--------- 2 files changed, 30 insertions(+), 37 deletions(-) diff --git a/src/promptflow/promptflow/executor/_line_execution_process_pool.py b/src/promptflow/promptflow/executor/_line_execution_process_pool.py index be67dbd9d33..263ea6ce58e 100644 --- a/src/promptflow/promptflow/executor/_line_execution_process_pool.py +++ b/src/promptflow/promptflow/executor/_line_execution_process_pool.py @@ -151,7 +151,6 @@ def __init__( self._validate_inputs = validate_inputs self._worker_count = flow_executor._worker_count multiprocessing_start_method = os.environ.get("PF_BATCH_METHOD") - self._DEFAULT_WORKER_COUNT = flow_executor._DEFAULT_WORKER_COUNT self._use_default_worker_count = flow_executor._use_default_worker_count sys_start_methods = multiprocessing.get_all_start_methods() if multiprocessing_start_method and multiprocessing_start_method not in sys_start_methods: @@ -203,31 +202,31 @@ def __enter__(self): if self._use_default_worker_count: self._n_process = min(self._worker_count, self._nlines, estimated_available_worker_count) bulk_logger.info( - f"Not using fork, and the environment variable PF_WORKER_COUNT is not set. Calculate the current " - f"system available memory divided by process memory, and take the minimum value of this value: " - f"{estimated_available_worker_count}, the default value for worker_count: " - f"{self._DEFAULT_WORKER_COUNT} and the row count: {self._nlines} as the final number of processes. " - f"process count: {self._n_process}") + f"Not using fork, the environment variable PF_WORKER_COUNT is not set. Calculate the estimated " + f"available worker count based on the current system available memory and process memory. Take the " + f"minimum value among the calculated value: {estimated_available_worker_count}, " + f"the default value for worker_count: {self._worker_count}, and the row count: " + f"{self._nlines}. Process count: {self._n_process}") else: self._n_process = self._worker_count if estimated_available_worker_count < self._worker_count: bulk_logger.warning( - f"The maximum number of processes calculated based on the system available memory " - f"is {estimated_available_worker_count}, and the PF_WORKER_COUNT is set to " - f"{self._worker_count}. Use the PF_WORKER_COUNT:{self._worker_count} as the final " - f"number of processes. process count: {self._n_process}") + f"The estimated available worker count calculated based on the system available memory " + f"is {estimated_available_worker_count}, but the PF_WORKER_COUNT is set to " + f"{self._worker_count}. This may affect optimal memory usage and performance. " + f"Process count: {self._n_process}") else: - log_process_count_info(self._worker_count, self._n_process) + log_process_count_info(self._worker_count) else: if self._use_default_worker_count: self._n_process = min(self._worker_count, self._nlines) bulk_logger.info( - f"Using fork, and the environment variable PF_WORKER_COUNT is not set. The number of processes is " - f"determined by the lesser of the default value for worker_count {self._DEFAULT_WORKER_COUNT} and " - f"the row count: {self._nlines}. process count: {self._n_process}") + f"Using fork, the environment variable PF_WORKER_COUNT is not set. Take the " + f"minimum value among the default value for worker_count {self._worker_count} and " + f"the row count: {self._nlines}. Process count: {self._n_process}") else: self._n_process = self._worker_count - log_process_count_info(self._worker_count, self._n_process) + log_process_count_info(self._worker_count) pool = ThreadPool(self._n_process, initializer=set_context, initargs=(contextvars.copy_context(),)) self._pool = pool @@ -597,8 +596,6 @@ def get_multiprocessing_context(multiprocessing_start_method=None): return context -def log_process_count_info(worker_count, process_count): +def log_process_count_info(worker_count): bulk_logger.info( - f"The number of processes has been explicitly set to {worker_count} by the " - f"environment variable PF_WORKER_COUNT. This value will determine the degree of parallelism of " - f"the process. process count: {process_count}") + f"Process count set to {worker_count} based on 'PF_WORKER_COUNT' environment variable.") diff --git a/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py b/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py index 5108b956772..70758b7fdef 100644 --- a/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py +++ b/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py @@ -82,14 +82,12 @@ def test_fork_mode_parallelism_in_subprocess( assert pool._n_process == n_process if is_set_environ_pf_worker_count: mock_logger.info.assert_called_with( - f"The number of processes has been explicitly set to {pf_worker_count} by the " - f"environment variable PF_WORKER_COUNT. This value will determine the degree of " - f"parallelism of the process. process count: {n_process}") + f"Process count set to {pf_worker_count} based on 'PF_WORKER_COUNT' environment variable.") else: mock_logger.info.assert_called_with( - f"Using fork, and the environment variable PF_WORKER_COUNT is not set. The number of processes " - f"is determined by the lesser of the default value for worker_count " - f"{pool._DEFAULT_WORKER_COUNT} and the row count: {nlines}. process count: {n_process}") + f"Using fork, the environment variable PF_WORKER_COUNT is not set. Take the " + f"minimum value among the default value for worker_count {pool._worker_count} and " + f"the row count: {nlines}. Process count: {n_process}") @pytest.mark.skip("This is a subprocess function used for testing and cannot be tested alone.") @@ -129,22 +127,20 @@ def test_spawn_mode_parallelism_in_subprocess( assert pool._n_process == n_process if is_set_environ_pf_worker_count and is_calculation_smaller_than_set: mock_logger.warning.assert_called_with( - f"The maximum number of processes calculated based on the system available memory " - f"is {available_max_worker_count}, and the PF_WORKER_COUNT is set to {pf_worker_count}" - f". Use the PF_WORKER_COUNT:{pf_worker_count} as the final number of processes. " - f"process count: {n_process}") + f"The estimated available worker count calculated based on the system available memory " + f"is {available_max_worker_count}, but the PF_WORKER_COUNT is set to " + f"{pf_worker_count}. This may affect optimal memory usage and performance. " + f"Process count: {n_process}") elif is_set_environ_pf_worker_count and not is_calculation_smaller_than_set: mock_logger.info.assert_called_with( - f"The number of processes has been explicitly set to {pf_worker_count} by the " - f"environment variable PF_WORKER_COUNT. This value will determine the degree of " - f"parallelism of the process. process count: {n_process}") + f"Process count set to {pf_worker_count} based on 'PF_WORKER_COUNT' environment variable.") elif not is_set_environ_pf_worker_count: mock_logger.info.assert_called_with( - f"Not using fork, and the environment variable PF_WORKER_COUNT is not set. Calculate " - f"the current system available memory divided by process memory, and take the minimum " - f"value of this value: {available_max_worker_count}, the default value for worker_count" - f": {pool._DEFAULT_WORKER_COUNT} and the row count: {nlines} as the final number of " - f"processes. process count: {n_process}") + f"Not using fork, the environment variable PF_WORKER_COUNT is not set. Calculate the " + f"estimated available worker count based on the current system available memory and " + f"process memory. Take the minimum value among the calculated value: " + f"{available_max_worker_count}, the default value for worker_count: {pool._worker_count}, " + f"and the row count: {nlines}. Process count: {n_process}") @pytest.mark.unittest From 7273b638a0873f663655d0a1cc75c673fcddbd2e Mon Sep 17 00:00:00 2001 From: LuLu Zuo Date: Tue, 5 Dec 2023 15:52:41 +0800 Subject: [PATCH 17/32] fix comments --- .../executor/_line_execution_process_pool.py | 107 +++++++++++------- .../promptflow/executor/flow_executor.py | 31 ----- .../unittests/executor/test_flow_executor.py | 33 +----- .../test_line_execution_process_pool.py | 105 ++++++++++------- 4 files changed, 133 insertions(+), 143 deletions(-) diff --git a/src/promptflow/promptflow/executor/_line_execution_process_pool.py b/src/promptflow/promptflow/executor/_line_execution_process_pool.py index 263ea6ce58e..c1649f870ca 100644 --- a/src/promptflow/promptflow/executor/_line_execution_process_pool.py +++ b/src/promptflow/promptflow/executor/_line_execution_process_pool.py @@ -1,5 +1,4 @@ import contextvars -import math import multiprocessing import os import queue @@ -136,6 +135,8 @@ def format_current_process(self, line_number: int, is_completed=False): class LineExecutionProcessPool: + _DEFAULT_WORKER_COUNT = 16 + def __init__( self, flow_executor: FlowExecutor, @@ -149,9 +150,10 @@ def __init__( self._run_id = run_id self._variant_id = variant_id self._validate_inputs = validate_inputs - self._worker_count = flow_executor._worker_count + # _worker_count: The number of workers to use for parallel execution of the Flow. + # _use_default_worker_count: Whether to use the default worker count. + self._use_default_worker_count, self._worker_count = self.load_worker_count_in_env() multiprocessing_start_method = os.environ.get("PF_BATCH_METHOD") - self._use_default_worker_count = flow_executor._use_default_worker_count sys_start_methods = multiprocessing.get_all_start_methods() if multiprocessing_start_method and multiprocessing_start_method not in sys_start_methods: bulk_logger.warning( @@ -197,36 +199,12 @@ def __enter__(self): self._inputs_queue = Queue() # Starting a new process in non-fork mode requires to allocate memory. Determine the maximum number of processes # based on available memory to avoid memory bursting. - if not self._use_fork: - estimated_available_worker_count = get_available_max_worker_count() - if self._use_default_worker_count: - self._n_process = min(self._worker_count, self._nlines, estimated_available_worker_count) - bulk_logger.info( - f"Not using fork, the environment variable PF_WORKER_COUNT is not set. Calculate the estimated " - f"available worker count based on the current system available memory and process memory. Take the " - f"minimum value among the calculated value: {estimated_available_worker_count}, " - f"the default value for worker_count: {self._worker_count}, and the row count: " - f"{self._nlines}. Process count: {self._n_process}") - else: - self._n_process = self._worker_count - if estimated_available_worker_count < self._worker_count: - bulk_logger.warning( - f"The estimated available worker count calculated based on the system available memory " - f"is {estimated_available_worker_count}, but the PF_WORKER_COUNT is set to " - f"{self._worker_count}. This may affect optimal memory usage and performance. " - f"Process count: {self._n_process}") - else: - log_process_count_info(self._worker_count) + if self._use_default_worker_count: + self._n_process = self._determine_worker_count() else: - if self._use_default_worker_count: - self._n_process = min(self._worker_count, self._nlines) - bulk_logger.info( - f"Using fork, the environment variable PF_WORKER_COUNT is not set. Take the " - f"minimum value among the default value for worker_count {self._worker_count} and " - f"the row count: {self._nlines}. Process count: {self._n_process}") - else: - self._n_process = self._worker_count - log_process_count_info(self._worker_count) + self._n_process = self._worker_count + self._log_process_count_info() + pool = ThreadPool(self._n_process, initializer=set_context, initargs=(contextvars.copy_context(),)) self._pool = pool @@ -435,6 +413,60 @@ def _generate_thread_status_messages(self, pool: ThreadPool, total_count: int): msgs.append("Processing Lines: " + ", ".join(lines) + ".") return msgs + def _determine_worker_count(self): + if not self._use_fork: + estimated_available_worker_count = get_available_max_worker_count() + n_process = min(self._worker_count, self._nlines, estimated_available_worker_count) + bulk_logger.info("Not using fork to create new process") + bulk_logger.info( + "The environment variable PF_WORKER_COUNT is not set or invalid. Calculate the worker count based on " + "the currently memory usage") + bulk_logger.info( + f"Calculated process count ({n_process}) by taking the minimum value among estimated process " + f"count ({estimated_available_worker_count}), the row count ({self._nlines}) and the default " + f"worker count ({self._worker_count}).") + return n_process + else: + n_process = min(self._worker_count, self._nlines) + bulk_logger.info("Using fork to create new process") + bulk_logger.info( + f"Calculated process count ({n_process}) by taking the minimum value among the the " + f"default value for worker_count ({self._worker_count}) and the row count ({self._nlines})") + return n_process + + def load_worker_count_in_env(self): + try: + pf_worker_count = os.environ.get("PF_WORKER_COUNT") + if pf_worker_count is None: + use_default_worker_count = True + worker_count = self._DEFAULT_WORKER_COUNT + else: + use_default_worker_count = False + worker_count = int(pf_worker_count) + except Exception as e: + bulk_logger.warning(f"Failed to convert PF_WORKER_COUNT '{pf_worker_count}' to an integer: {e}") + use_default_worker_count = True + worker_count = self._DEFAULT_WORKER_COUNT + + if worker_count <= 0: + bulk_logger.warning( + f"Invalid worker count: {worker_count}. Resetting to default value: {self._DEFAULT_WORKER_COUNT}") + use_default_worker_count = True + worker_count = self._DEFAULT_WORKER_COUNT + + return use_default_worker_count, worker_count + + def _log_process_count_info(self): + bulk_logger.info( + f"Process count set to {self._worker_count} based on 'PF_WORKER_COUNT' environment variable.") + if not self._use_fork: + estimated_available_worker_count = get_available_max_worker_count() + if estimated_available_worker_count < self._worker_count: + bulk_logger.warning( + f"The estimated available worker count calculated based on the system available memory " + f"is {estimated_available_worker_count}, but the PF_WORKER_COUNT is set to " + f"{self._worker_count}. This may affect optimal memory usage and performance. ") + def _exec_line( executor: FlowExecutor, @@ -568,7 +600,7 @@ def get_available_max_worker_count(): process = psutil.Process(pid) process_memory_info = process.memory_info() process_memory = process_memory_info.rss / (1024 * 1024) # in MB - estimated_available_worker_count = math.floor((available_memory) / process_memory) + estimated_available_worker_count = available_memory // process_memory if estimated_available_worker_count < 1: # TODO: For the case of vector db, Optimize execution logic # 1. Let the main process not consume memory because it does not actually invoke @@ -578,8 +610,8 @@ def get_available_max_worker_count(): f"Available max worker count {estimated_available_worker_count} is less than 1, set it to 1.") estimated_available_worker_count = 1 bulk_logger.info( - f"Process {pid} current available memory is {process_memory}, " - f"memory consumption of current process is {available_memory}, " + f"Current system's available memory is {available_memory}MB, " + f"memory consumption of current process is {process_memory}MB, " f"estimated available worker count is {available_memory}/{process_memory} " f"= {estimated_available_worker_count}" ) @@ -594,8 +626,3 @@ def get_multiprocessing_context(multiprocessing_start_method=None): else: context = multiprocessing.get_context() return context - - -def log_process_count_info(worker_count): - bulk_logger.info( - f"Process count set to {worker_count} based on 'PF_WORKER_COUNT' environment variable.") diff --git a/src/promptflow/promptflow/executor/flow_executor.py b/src/promptflow/promptflow/executor/flow_executor.py index e4b8ecada88..375f7bf8842 100644 --- a/src/promptflow/promptflow/executor/flow_executor.py +++ b/src/promptflow/promptflow/executor/flow_executor.py @@ -5,7 +5,6 @@ import copy import functools import inspect -import os import uuid from pathlib import Path from threading import current_thread @@ -88,7 +87,6 @@ def __init__( cache_manager: AbstractCacheManager, loaded_tools: Mapping[str, Callable], *, - worker_count=None, raise_ex: bool = False, working_dir=None, line_timeout_sec=LINE_TIMEOUT_SEC, @@ -106,8 +104,6 @@ def __init__( :type cache_manager: ~promptflow._core.cache_manager.AbstractCacheManager :param loaded_tools: A mapping of tool names to their corresponding functions. :type loaded_tools: Mapping[str, Callable] - :param worker_count: The number of workers to use for parallel execution of the Flow. - :type worker_count: int or None :param raise_ex: Whether to raise an exception if an error occurs during execution. :type raise_ex: bool :param working_dir: The working directory to use for execution. @@ -126,10 +122,6 @@ def __init__( self._connections = connections self._aggregation_inputs_references = get_aggregation_inputs_properties(flow) self._aggregation_nodes = {node.name for node in self._flow.nodes if node.aggregation} - if worker_count is not None: - self._worker_count = worker_count - else: - self._use_default_worker_count, self._worker_count = load_worker_count_in_env(self._DEFAULT_WORKER_COUNT) self._run_tracker = run_tracker self._cache_manager = cache_manager self._loaded_tools = loaded_tools @@ -957,29 +949,6 @@ def ensure_flow_is_serializable(self): self._tools_manager.wrap_tool(node.name, wrapper=_ensure_node_result_is_serializable) -def load_worker_count_in_env(DEFAULT_WORKER_COUNT): - try: - pf_worker_count = os.environ.get("PF_WORKER_COUNT") - if pf_worker_count is None: - use_default_worker_count = True - worker_count = DEFAULT_WORKER_COUNT - else: - use_default_worker_count = False - worker_count = int(pf_worker_count) - except Exception as e: - logger.warning(f"Failed to convert PF_WORKER_COUNT '{pf_worker_count}' to an integer: {e}") - use_default_worker_count = True - worker_count = DEFAULT_WORKER_COUNT - - if worker_count <= 0: - logger.warning( - f"Invalid worker count: {worker_count}. Resetting to default value: {DEFAULT_WORKER_COUNT}") - use_default_worker_count = True - worker_count = DEFAULT_WORKER_COUNT - - return use_default_worker_count, worker_count - - def _inject_stream_options(should_stream: Callable[[], bool], streaming_option_parameter="stream"): """Inject the stream options to the decorated function. diff --git a/src/promptflow/tests/executor/unittests/executor/test_flow_executor.py b/src/promptflow/tests/executor/unittests/executor/test_flow_executor.py index a2a5988c55d..784ca44e6ef 100644 --- a/src/promptflow/tests/executor/unittests/executor/test_flow_executor.py +++ b/src/promptflow/tests/executor/unittests/executor/test_flow_executor.py @@ -1,7 +1,6 @@ -from unittest.mock import Mock, patch +from unittest.mock import Mock import pytest -import os from promptflow import tool from promptflow.contracts.flow import FlowInputDefinition @@ -12,9 +11,7 @@ _ensure_node_result_is_serializable, _inject_stream_options, enable_streaming_for_llm_tool, - load_worker_count_in_env ) -from promptflow._utils.logger_utils import logger from promptflow.tools.aoai import chat, completion from promptflow.tools.embedding import embedding @@ -181,31 +178,3 @@ def test_streaming_tool_should_be_consumed_and_merged(self): def test_non_streaming_tool_should_not_be_affected(self): func = _ensure_node_result_is_serializable(non_streaming_tool) assert func() == 1 - - -class TestLoadWorkerCountInEnv: - def test_no_env_var_set(self): - with patch.dict(os.environ, {}, clear=True): - use_default, worker_count = load_worker_count_in_env(4) - assert use_default is True - assert worker_count == 4 - - def test_valid_integer_env_var(self): - with patch.dict(os.environ, {'PF_WORKER_COUNT': '5'}): - use_default, worker_count = load_worker_count_in_env(4) - assert use_default is False - assert worker_count == 5 - - def test_invalid_non_integer_env_var(self): - with patch.dict(os.environ, {'PF_WORKER_COUNT': 'invalid'}), patch.object(logger, 'warning') as mock_warning: - use_default, worker_count = load_worker_count_in_env(4) - mock_warning.assert_called() - assert use_default is True - assert worker_count == 4 - - def test_invalid_non_positive_integer_env_var(self): - with patch.dict(os.environ, {'PF_WORKER_COUNT': '0'}), patch.object(logger, 'warning') as mock_warning: - use_default, worker_count = load_worker_count_in_env(4) - mock_warning.assert_called() - assert use_default is True - assert worker_count == 4 diff --git a/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py b/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py index 70758b7fdef..c4274744b64 100644 --- a/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py +++ b/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py @@ -9,7 +9,7 @@ import pytest from pytest_mock import MockFixture -from promptflow._utils.logger_utils import LogContext +from promptflow._utils.logger_utils import LogContext, bulk_logger from promptflow.contracts.run_info import Status from promptflow.exceptions import ErrorTarget, UserErrorException from promptflow.executor import FlowExecutor @@ -81,13 +81,14 @@ def test_fork_mode_parallelism_in_subprocess( ) as pool: assert pool._n_process == n_process if is_set_environ_pf_worker_count: - mock_logger.info.assert_called_with( + mock_logger.info.assert_any_call( f"Process count set to {pf_worker_count} based on 'PF_WORKER_COUNT' environment variable.") else: - mock_logger.info.assert_called_with( - f"Using fork, the environment variable PF_WORKER_COUNT is not set. Take the " - f"minimum value among the default value for worker_count {pool._worker_count} and " - f"the row count: {nlines}. Process count: {n_process}") + mock_logger.info.assert_any_call("Using fork to create new process") + mock_logger.info.assert_any_call( + f"Calculated process count ({n_process}) by taking the minimum value among the the " + f"default value for worker_count ({pool._worker_count}) and the row count ({nlines})" + ) @pytest.mark.skip("This is a subprocess function used for testing and cannot be tested alone.") @@ -97,7 +98,7 @@ def test_spawn_mode_parallelism_in_subprocess( is_set_environ_pf_worker_count, is_calculation_smaller_than_set, pf_worker_count, - available_max_worker_count, + estimated_available_worker_count, n_process ): os.environ["PF_BATCH_METHOD"] = "spawn" @@ -126,26 +127,30 @@ def test_spawn_mode_parallelism_in_subprocess( ) as pool: assert pool._n_process == n_process if is_set_environ_pf_worker_count and is_calculation_smaller_than_set: - mock_logger.warning.assert_called_with( + mock_logger.info.assert_any_call( + f"Process count set to {pf_worker_count} based on 'PF_WORKER_COUNT' environment variable.") + mock_logger.warning.assert_any_call( f"The estimated available worker count calculated based on the system available memory " - f"is {available_max_worker_count}, but the PF_WORKER_COUNT is set to " - f"{pf_worker_count}. This may affect optimal memory usage and performance. " - f"Process count: {n_process}") + f"is {estimated_available_worker_count}, but the PF_WORKER_COUNT is set to " + f"{pf_worker_count}. This may affect optimal memory usage and performance. ") elif is_set_environ_pf_worker_count and not is_calculation_smaller_than_set: - mock_logger.info.assert_called_with( + mock_logger.info.assert_any_call( f"Process count set to {pf_worker_count} based on 'PF_WORKER_COUNT' environment variable.") elif not is_set_environ_pf_worker_count: - mock_logger.info.assert_called_with( - f"Not using fork, the environment variable PF_WORKER_COUNT is not set. Calculate the " - f"estimated available worker count based on the current system available memory and " - f"process memory. Take the minimum value among the calculated value: " - f"{available_max_worker_count}, the default value for worker_count: {pool._worker_count}, " - f"and the row count: {nlines}. Process count: {n_process}") + mock_logger.info.assert_any_call("Not using fork to create new process") + mock_logger.info.assert_any_call( + "The environment variable PF_WORKER_COUNT is not set or invalid. Calculate the worker " + "count based on the currently memory usage" + ) + mock_logger.info.assert_any_call( + f"Calculated process count ({n_process}) by taking the minimum value among estimated " + f"process count ({estimated_available_worker_count}), the row count ({nlines}) and the " + f"default worker count ({pool._worker_count})." + ) @pytest.mark.unittest class TestLineExecutionProcessPool: - def create_line_execution_process_pool(self, dev_connections): executor = FlowExecutor.create( get_yaml_file(SAMPLE_FLOW), @@ -373,14 +378,16 @@ def test_process_pool_parallelism_in_fork_mode( is_set_environ_pf_worker_count, pf_worker_count, n_process): - p = multiprocessing.Process(target=test_fork_mode_parallelism_in_subprocess, - args=(dev_connections, - flow_folder, - is_set_environ_pf_worker_count, - pf_worker_count, - n_process)) + p = multiprocessing.Process( + target=test_fork_mode_parallelism_in_subprocess, + args=(dev_connections, + flow_folder, + is_set_environ_pf_worker_count, + pf_worker_count, + n_process)) p.start() p.join() + assert p.exitcode == 0 @pytest.mark.parametrize( ( @@ -388,7 +395,7 @@ def test_process_pool_parallelism_in_fork_mode( "is_set_environ_pf_worker_count", "is_calculation_smaller_than_set", "pf_worker_count", - "available_max_worker_count", + "estimated_available_worker_count", "n_process" ), [ @@ -404,27 +411,45 @@ def test_process_pool_parallelism_in_non_spawn_mode( is_set_environ_pf_worker_count, is_calculation_smaller_than_set, pf_worker_count, - available_max_worker_count, + estimated_available_worker_count, n_process ): - p = multiprocessing.Process(target=test_spawn_mode_parallelism_in_subprocess, - args=(dev_connections, - flow_folder, - is_set_environ_pf_worker_count, - is_calculation_smaller_than_set, - pf_worker_count, - available_max_worker_count, - n_process)) + p = multiprocessing.Process( + target=test_spawn_mode_parallelism_in_subprocess, + args=(dev_connections, + flow_folder, + is_set_environ_pf_worker_count, + is_calculation_smaller_than_set, + pf_worker_count, + estimated_available_worker_count, + n_process)) p.start() p.join() + assert p.exitcode == 0 + + @pytest.mark.parametrize("env_var, expected_use_default, expected_worker_count", [ + ({}, True, 16), + ({'PF_WORKER_COUNT': '5'}, False, 5), + ({'PF_WORKER_COUNT': 'invalid'}, True, 16), + ({'PF_WORKER_COUNT': '0'}, True, 16) + ]) + def test_worker_count_setting(self, dev_connections, env_var, expected_use_default, expected_worker_count): + with patch.dict(os.environ, env_var), patch.object(bulk_logger, 'warning') as mock_warning: + line_execution_process_pool = self.create_line_execution_process_pool(dev_connections) + + if 'invalid' in env_var.get('PF_WORKER_COUNT', '') or '0' in env_var.get('PF_WORKER_COUNT', ''): + mock_warning.assert_called() + + assert line_execution_process_pool._use_default_worker_count == expected_use_default + assert line_execution_process_pool._worker_count == expected_worker_count class TestGetAvailableMaxWorkerCount: @pytest.mark.parametrize( "available_memory, process_memory, expected_max_worker_count, actual_calculate_worker_count", [ - (128.0, 64.0, 2, 2), # available_memory > 0 - (0.0, 64.0, 1, 0), # available_memory = 0 + (128.0, 64.0, 2, 2.0), # available_memory/process_memory > 1 + (63.0, 64.0, 1, 1), # available_memory/process_memory < 1 ], ) def test_get_available_max_worker_count( @@ -444,8 +469,8 @@ def test_get_available_max_worker_count( "set it to 1." ) mock_logger.info.assert_called_with( - f"Process {os.getpid()} current available memory is {process_memory}, " - f"memory consumption of current process is {available_memory}, " + f"Current system's available memory is {available_memory}MB, " + f"memory consumption of current process is {process_memory}MB, " f"estimated available worker count is {available_memory}/{process_memory} " - f"= {estimated_available_worker_count}" + f"= {actual_calculate_worker_count}" ) From 5d3f7f8c814e449d9a7a045a7f482db79f0b891b Mon Sep 17 00:00:00 2001 From: LuLu Zuo Date: Tue, 5 Dec 2023 16:27:26 +0800 Subject: [PATCH 18/32] Rename variable --- .../executor/_line_execution_process_pool.py | 22 +++++++++---------- .../test_line_execution_process_pool.py | 16 +++++++------- 2 files changed, 19 insertions(+), 19 deletions(-) diff --git a/src/promptflow/promptflow/executor/_line_execution_process_pool.py b/src/promptflow/promptflow/executor/_line_execution_process_pool.py index c1649f870ca..9b30a92669c 100644 --- a/src/promptflow/promptflow/executor/_line_execution_process_pool.py +++ b/src/promptflow/promptflow/executor/_line_execution_process_pool.py @@ -151,8 +151,8 @@ def __init__( self._variant_id = variant_id self._validate_inputs = validate_inputs # _worker_count: The number of workers to use for parallel execution of the Flow. - # _use_default_worker_count: Whether to use the default worker count. - self._use_default_worker_count, self._worker_count = self.load_worker_count_in_env() + # _is_pf_worker_count_set_and_valid: Whether to use the default worker count. + self._is_pf_worker_count_set_and_valid, self._worker_count = self.load_worker_count_in_env() multiprocessing_start_method = os.environ.get("PF_BATCH_METHOD") sys_start_methods = multiprocessing.get_all_start_methods() if multiprocessing_start_method and multiprocessing_start_method not in sys_start_methods: @@ -199,7 +199,7 @@ def __enter__(self): self._inputs_queue = Queue() # Starting a new process in non-fork mode requires to allocate memory. Determine the maximum number of processes # based on available memory to avoid memory bursting. - if self._use_default_worker_count: + if not self._is_pf_worker_count_set_and_valid: self._n_process = self._determine_worker_count() else: self._n_process = self._worker_count @@ -430,31 +430,31 @@ def _determine_worker_count(self): n_process = min(self._worker_count, self._nlines) bulk_logger.info("Using fork to create new process") bulk_logger.info( - f"Calculated process count ({n_process}) by taking the minimum value among the the " - f"default value for worker_count ({self._worker_count}) and the row count ({self._nlines})") + f"Calculated process count ({n_process}) by taking the minimum value among the " + f"default worker_count ({self._worker_count}) and the row count ({self._nlines})") return n_process def load_worker_count_in_env(self): try: pf_worker_count = os.environ.get("PF_WORKER_COUNT") if pf_worker_count is None: - use_default_worker_count = True + is_pf_worker_count_set_and_valid = False worker_count = self._DEFAULT_WORKER_COUNT else: - use_default_worker_count = False + is_pf_worker_count_set_and_valid = True worker_count = int(pf_worker_count) except Exception as e: bulk_logger.warning(f"Failed to convert PF_WORKER_COUNT '{pf_worker_count}' to an integer: {e}") - use_default_worker_count = True + is_pf_worker_count_set_and_valid = False worker_count = self._DEFAULT_WORKER_COUNT if worker_count <= 0: bulk_logger.warning( f"Invalid worker count: {worker_count}. Resetting to default value: {self._DEFAULT_WORKER_COUNT}") - use_default_worker_count = True + is_pf_worker_count_set_and_valid = False worker_count = self._DEFAULT_WORKER_COUNT - return use_default_worker_count, worker_count + return is_pf_worker_count_set_and_valid, worker_count def _log_process_count_info(self): bulk_logger.info( @@ -600,7 +600,7 @@ def get_available_max_worker_count(): process = psutil.Process(pid) process_memory_info = process.memory_info() process_memory = process_memory_info.rss / (1024 * 1024) # in MB - estimated_available_worker_count = available_memory // process_memory + estimated_available_worker_count = int(available_memory // process_memory) if estimated_available_worker_count < 1: # TODO: For the case of vector db, Optimize execution logic # 1. Let the main process not consume memory because it does not actually invoke diff --git a/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py b/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py index c4274744b64..10ae6995c48 100644 --- a/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py +++ b/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py @@ -86,8 +86,8 @@ def test_fork_mode_parallelism_in_subprocess( else: mock_logger.info.assert_any_call("Using fork to create new process") mock_logger.info.assert_any_call( - f"Calculated process count ({n_process}) by taking the minimum value among the the " - f"default value for worker_count ({pool._worker_count}) and the row count ({nlines})" + f"Calculated process count ({n_process}) by taking the minimum value among the " + f"default worker_count ({pool._worker_count}) and the row count ({nlines})" ) @@ -428,10 +428,10 @@ def test_process_pool_parallelism_in_non_spawn_mode( assert p.exitcode == 0 @pytest.mark.parametrize("env_var, expected_use_default, expected_worker_count", [ - ({}, True, 16), - ({'PF_WORKER_COUNT': '5'}, False, 5), - ({'PF_WORKER_COUNT': 'invalid'}, True, 16), - ({'PF_WORKER_COUNT': '0'}, True, 16) + ({}, False, 16), + ({'PF_WORKER_COUNT': '5'}, True, 5), + ({'PF_WORKER_COUNT': 'invalid'}, False, 16), + ({'PF_WORKER_COUNT': '0'}, False, 16) ]) def test_worker_count_setting(self, dev_connections, env_var, expected_use_default, expected_worker_count): with patch.dict(os.environ, env_var), patch.object(bulk_logger, 'warning') as mock_warning: @@ -440,7 +440,7 @@ def test_worker_count_setting(self, dev_connections, env_var, expected_use_defau if 'invalid' in env_var.get('PF_WORKER_COUNT', '') or '0' in env_var.get('PF_WORKER_COUNT', ''): mock_warning.assert_called() - assert line_execution_process_pool._use_default_worker_count == expected_use_default + assert line_execution_process_pool._is_pf_worker_count_set_and_valid == expected_use_default assert line_execution_process_pool._worker_count == expected_worker_count @@ -448,7 +448,7 @@ class TestGetAvailableMaxWorkerCount: @pytest.mark.parametrize( "available_memory, process_memory, expected_max_worker_count, actual_calculate_worker_count", [ - (128.0, 64.0, 2, 2.0), # available_memory/process_memory > 1 + (128.0, 64.0, 2, 2), # available_memory/process_memory > 1 (63.0, 64.0, 1, 1), # available_memory/process_memory < 1 ], ) From 0bed290c0bb78eec5ffdbd39e3289e3137a3618a Mon Sep 17 00:00:00 2001 From: LuLu Zuo Date: Tue, 5 Dec 2023 18:42:54 +0800 Subject: [PATCH 19/32] fix comments --- src/promptflow/promptflow/_utils/utils.py | 18 ++++- .../executor/_line_execution_process_pool.py | 78 +++++++------------ .../executor/unittests/_utils/test_utils.py | 28 ++++++- .../test_line_execution_process_pool.py | 22 +----- 4 files changed, 76 insertions(+), 70 deletions(-) diff --git a/src/promptflow/promptflow/_utils/utils.py b/src/promptflow/promptflow/_utils/utils.py index fe7b3543906..5ae9dcd5f1c 100644 --- a/src/promptflow/promptflow/_utils/utils.py +++ b/src/promptflow/promptflow/_utils/utils.py @@ -167,7 +167,7 @@ def extract_user_frame_summaries(frame_summaries: List[traceback.FrameSummary]): # If the current frame is in tool.py and the next frame is not in _core folder # then we can say that the next frame is in user code. if cur_file == tool_file and not next_file.startswith(core_folder): - return frame_summaries[i + 1 :] + return frame_summaries[i + 1:] return frame_summaries @@ -247,3 +247,19 @@ def parse_ua_to_dict(ua): key, value = item.split("/") ua_dict[key] = value return ua_dict + + +def get_int_env_var(env_var_name, default_value=None): + """ + The function `get_int_env_var` retrieves an integer environment variable value, with an optional + default value if the variable is not set or cannot be converted to an integer. + + :param env_var_name: The name of the environment variable you want to retrieve the value of + :param default_value: The default value is the value that will be returned if the environment + variable is not found or if it cannot be converted to an integer + :return: an integer value. + """ + try: + return int(os.environ.get(env_var_name, default_value)) + except Exception: + return default_value diff --git a/src/promptflow/promptflow/executor/_line_execution_process_pool.py b/src/promptflow/promptflow/executor/_line_execution_process_pool.py index 9b30a92669c..5c6d775114b 100644 --- a/src/promptflow/promptflow/executor/_line_execution_process_pool.py +++ b/src/promptflow/promptflow/executor/_line_execution_process_pool.py @@ -20,7 +20,7 @@ from promptflow._utils.logger_utils import LogContext, bulk_logger from promptflow._utils.multimedia_utils import _process_recursively, persist_multimedia_data from promptflow._utils.thread_utils import RepeatLogTimer -from promptflow._utils.utils import log_progress, set_context +from promptflow._utils.utils import log_progress, set_context, get_int_env_var from promptflow.contracts.multimedia import Image from promptflow.contracts.run_info import FlowRunInfo from promptflow.contracts.run_info import RunInfo as NodeRunInfo @@ -150,9 +150,6 @@ def __init__( self._run_id = run_id self._variant_id = variant_id self._validate_inputs = validate_inputs - # _worker_count: The number of workers to use for parallel execution of the Flow. - # _is_pf_worker_count_set_and_valid: Whether to use the default worker count. - self._is_pf_worker_count_set_and_valid, self._worker_count = self.load_worker_count_in_env() multiprocessing_start_method = os.environ.get("PF_BATCH_METHOD") sys_start_methods = multiprocessing.get_all_start_methods() if multiprocessing_start_method and multiprocessing_start_method not in sys_start_methods: @@ -197,13 +194,7 @@ def __enter__(self): self._completed_idx = manager.dict() self._inputs_queue = Queue() - # Starting a new process in non-fork mode requires to allocate memory. Determine the maximum number of processes - # based on available memory to avoid memory bursting. - if not self._is_pf_worker_count_set_and_valid: - self._n_process = self._determine_worker_count() - else: - self._n_process = self._worker_count - self._log_process_count_info() + self._n_process = self._determine_worker_count() pool = ThreadPool(self._n_process, initializer=set_context, initargs=(contextvars.copy_context(),)) self._pool = pool @@ -414,47 +405,36 @@ def _generate_thread_status_messages(self, pool: ThreadPool, total_count: int): return msgs def _determine_worker_count(self): - if not self._use_fork: - estimated_available_worker_count = get_available_max_worker_count() - n_process = min(self._worker_count, self._nlines, estimated_available_worker_count) - bulk_logger.info("Not using fork to create new process") + worker_count = get_int_env_var("PF_WORKER_COUNT") + estimated_available_worker_count = get_available_max_worker_count() if not self._use_fork else None + if worker_count is not None and worker_count > 0: bulk_logger.info( - "The environment variable PF_WORKER_COUNT is not set or invalid. Calculate the worker count based on " - "the currently memory usage") - bulk_logger.info( - f"Calculated process count ({n_process}) by taking the minimum value among estimated process " - f"count ({estimated_available_worker_count}), the row count ({self._nlines}) and the default " - f"worker count ({self._worker_count}).") - return n_process + f"Process count set to {worker_count} based on 'PF_WORKER_COUNT' environment variable.") + if not self._use_fork and estimated_available_worker_count < worker_count: + bulk_logger.warning( + f"The estimated available worker count calculated based on the system available memory " + f"is {estimated_available_worker_count}, but the PF_WORKER_COUNT is set to " + f"{worker_count}. This may affect optimal memory usage and performance. ") + return worker_count else: - n_process = min(self._worker_count, self._nlines) - bulk_logger.info("Using fork to create new process") - bulk_logger.info( - f"Calculated process count ({n_process}) by taking the minimum value among the " - f"default worker_count ({self._worker_count}) and the row count ({self._nlines})") - return n_process - - def load_worker_count_in_env(self): - try: - pf_worker_count = os.environ.get("PF_WORKER_COUNT") - if pf_worker_count is None: - is_pf_worker_count_set_and_valid = False - worker_count = self._DEFAULT_WORKER_COUNT + worker_count = min(self._DEFAULT_WORKER_COUNT, self._nlines) + if not self._use_fork: + worker_count = min(worker_count, estimated_available_worker_count) + bulk_logger.info("Not using fork to create new process") + bulk_logger.info( + "The environment variable PF_WORKER_COUNT is not set or invalid. Calculate the worker count based " + "on the currently memory usage") + bulk_logger.info( + f"Calculated process count ({worker_count}) by taking the minimum value among estimated process " + f"count ({estimated_available_worker_count}), the row count ({self._nlines}) and the default " + f"worker count ({self._DEFAULT_WORKER_COUNT}).") + return worker_count else: - is_pf_worker_count_set_and_valid = True - worker_count = int(pf_worker_count) - except Exception as e: - bulk_logger.warning(f"Failed to convert PF_WORKER_COUNT '{pf_worker_count}' to an integer: {e}") - is_pf_worker_count_set_and_valid = False - worker_count = self._DEFAULT_WORKER_COUNT - - if worker_count <= 0: - bulk_logger.warning( - f"Invalid worker count: {worker_count}. Resetting to default value: {self._DEFAULT_WORKER_COUNT}") - is_pf_worker_count_set_and_valid = False - worker_count = self._DEFAULT_WORKER_COUNT - - return is_pf_worker_count_set_and_valid, worker_count + bulk_logger.info("Using fork to create new process") + bulk_logger.info( + f"Calculated process count ({worker_count}) by taking the minimum value among the " + f"default worker_count ({self._DEFAULT_WORKER_COUNT}) and the row count ({self._nlines})") + return worker_count def _log_process_count_info(self): bulk_logger.info( diff --git a/src/promptflow/tests/executor/unittests/_utils/test_utils.py b/src/promptflow/tests/executor/unittests/_utils/test_utils.py index d85ff4ff800..617d9eec5ef 100644 --- a/src/promptflow/tests/executor/unittests/_utils/test_utils.py +++ b/src/promptflow/tests/executor/unittests/_utils/test_utils.py @@ -1,6 +1,8 @@ import pytest +import os +from unittest.mock import patch -from promptflow._utils.utils import is_json_serializable +from promptflow._utils.utils import is_json_serializable, get_int_env_var class MyObj: @@ -12,3 +14,27 @@ class TestUtils: @pytest.mark.parametrize("value, expected_res", [(None, True), (1, True), ("", True), (MyObj(), False)]) def test_is_json_serializable(self, value, expected_res): assert is_json_serializable(value) == expected_res + + @pytest.mark.parametrize( + "env_var, env_value, default_value, expected_result", + [ + ("TEST_VAR", "10", None, 10), # Valid integer string + ("TEST_VAR", "invalid", None, None), # Invalid integer strings + ("TEST_VAR", None, 5, 5), # Environment variable does not exist + ("TEST_VAR", "10", 5, 10), # Valid integer string with a default value + ("TEST_VAR", "invalid", 5, 5), # Invalid integer string with a default value + ]) + def test_get_int_env_var(self, env_var, env_value, default_value, expected_result): + with patch.dict(os.environ, {env_var: env_value} if env_value is not None else {}): + assert get_int_env_var(env_var, default_value) == expected_result + + @pytest.mark.parametrize( + "env_var, env_value, expected_result", + [ + ("TEST_VAR", "10", 10), # Valid integer string + ("TEST_VAR", "invalid", None), # Invalid integer strings + ("TEST_VAR", None, None), # Environment variable does not exist + ]) + def test_get_int_env_var_without_default_vaue(self, env_var, env_value, expected_result): + with patch.dict(os.environ, {env_var: env_value} if env_value is not None else {}): + assert get_int_env_var(env_var) == expected_result diff --git a/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py b/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py index 10ae6995c48..1a44d92b724 100644 --- a/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py +++ b/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py @@ -9,7 +9,7 @@ import pytest from pytest_mock import MockFixture -from promptflow._utils.logger_utils import LogContext, bulk_logger +from promptflow._utils.logger_utils import LogContext from promptflow.contracts.run_info import Status from promptflow.exceptions import ErrorTarget, UserErrorException from promptflow.executor import FlowExecutor @@ -87,7 +87,7 @@ def test_fork_mode_parallelism_in_subprocess( mock_logger.info.assert_any_call("Using fork to create new process") mock_logger.info.assert_any_call( f"Calculated process count ({n_process}) by taking the minimum value among the " - f"default worker_count ({pool._worker_count}) and the row count ({nlines})" + f"default worker_count ({pool._DEFAULT_WORKER_COUNT}) and the row count ({nlines})" ) @@ -145,7 +145,7 @@ def test_spawn_mode_parallelism_in_subprocess( mock_logger.info.assert_any_call( f"Calculated process count ({n_process}) by taking the minimum value among estimated " f"process count ({estimated_available_worker_count}), the row count ({nlines}) and the " - f"default worker count ({pool._worker_count})." + f"default worker count ({pool._DEFAULT_WORKER_COUNT})." ) @@ -427,22 +427,6 @@ def test_process_pool_parallelism_in_non_spawn_mode( p.join() assert p.exitcode == 0 - @pytest.mark.parametrize("env_var, expected_use_default, expected_worker_count", [ - ({}, False, 16), - ({'PF_WORKER_COUNT': '5'}, True, 5), - ({'PF_WORKER_COUNT': 'invalid'}, False, 16), - ({'PF_WORKER_COUNT': '0'}, False, 16) - ]) - def test_worker_count_setting(self, dev_connections, env_var, expected_use_default, expected_worker_count): - with patch.dict(os.environ, env_var), patch.object(bulk_logger, 'warning') as mock_warning: - line_execution_process_pool = self.create_line_execution_process_pool(dev_connections) - - if 'invalid' in env_var.get('PF_WORKER_COUNT', '') or '0' in env_var.get('PF_WORKER_COUNT', ''): - mock_warning.assert_called() - - assert line_execution_process_pool._is_pf_worker_count_set_and_valid == expected_use_default - assert line_execution_process_pool._worker_count == expected_worker_count - class TestGetAvailableMaxWorkerCount: @pytest.mark.parametrize( From c6bd663b4e18f7a3ddcd1a34cad8c845f70f6dcd Mon Sep 17 00:00:00 2001 From: LuLu Zuo Date: Tue, 5 Dec 2023 18:45:19 +0800 Subject: [PATCH 20/32] delete useless function --- .../executor/_line_execution_process_pool.py | 11 ----------- 1 file changed, 11 deletions(-) diff --git a/src/promptflow/promptflow/executor/_line_execution_process_pool.py b/src/promptflow/promptflow/executor/_line_execution_process_pool.py index 5c6d775114b..a29cdf68d8b 100644 --- a/src/promptflow/promptflow/executor/_line_execution_process_pool.py +++ b/src/promptflow/promptflow/executor/_line_execution_process_pool.py @@ -436,17 +436,6 @@ def _determine_worker_count(self): f"default worker_count ({self._DEFAULT_WORKER_COUNT}) and the row count ({self._nlines})") return worker_count - def _log_process_count_info(self): - bulk_logger.info( - f"Process count set to {self._worker_count} based on 'PF_WORKER_COUNT' environment variable.") - if not self._use_fork: - estimated_available_worker_count = get_available_max_worker_count() - if estimated_available_worker_count < self._worker_count: - bulk_logger.warning( - f"The estimated available worker count calculated based on the system available memory " - f"is {estimated_available_worker_count}, but the PF_WORKER_COUNT is set to " - f"{self._worker_count}. This may affect optimal memory usage and performance. ") - def _exec_line( executor: FlowExecutor, From 76286b713017960551ea161f73aada08a61b3c62 Mon Sep 17 00:00:00 2001 From: LuLu Zuo Date: Tue, 5 Dec 2023 21:53:33 +0800 Subject: [PATCH 21/32] remove flow_executor's worker_count --- .../promptflow/executor/_line_execution_process_pool.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/promptflow/promptflow/executor/_line_execution_process_pool.py b/src/promptflow/promptflow/executor/_line_execution_process_pool.py index a29cdf68d8b..e26e10bbc49 100644 --- a/src/promptflow/promptflow/executor/_line_execution_process_pool.py +++ b/src/promptflow/promptflow/executor/_line_execution_process_pool.py @@ -506,7 +506,6 @@ def create_executor_fork(*, flow_executor: FlowExecutor, storage: AbstractRunSto run_tracker=run_tracker, cache_manager=flow_executor._cache_manager, loaded_tools=flow_executor._loaded_tools, - worker_count=flow_executor._worker_count, raise_ex=False, line_timeout_sec=flow_executor._line_timeout_sec, ) From f06266171d4c323ca186a9793041b2d447ec1220 Mon Sep 17 00:00:00 2001 From: LuLu Zuo Date: Tue, 5 Dec 2023 23:04:02 +0800 Subject: [PATCH 22/32] Optimize log content --- .../executor/_line_execution_process_pool.py | 10 +++++----- .../test_line_execution_process_pool.py | 16 +++++++++------- 2 files changed, 14 insertions(+), 12 deletions(-) diff --git a/src/promptflow/promptflow/executor/_line_execution_process_pool.py b/src/promptflow/promptflow/executor/_line_execution_process_pool.py index e26e10bbc49..39ff7c1dcaf 100644 --- a/src/promptflow/promptflow/executor/_line_execution_process_pool.py +++ b/src/promptflow/promptflow/executor/_line_execution_process_pool.py @@ -409,7 +409,7 @@ def _determine_worker_count(self): estimated_available_worker_count = get_available_max_worker_count() if not self._use_fork else None if worker_count is not None and worker_count > 0: bulk_logger.info( - f"Process count set to {worker_count} based on 'PF_WORKER_COUNT' environment variable.") + f"Process count set to {worker_count} based on the 'PF_WORKER_COUNT' environment variable.") if not self._use_fork and estimated_available_worker_count < worker_count: bulk_logger.warning( f"The estimated available worker count calculated based on the system available memory " @@ -420,20 +420,20 @@ def _determine_worker_count(self): worker_count = min(self._DEFAULT_WORKER_COUNT, self._nlines) if not self._use_fork: worker_count = min(worker_count, estimated_available_worker_count) - bulk_logger.info("Not using fork to create new process") + bulk_logger.info("Not using fork to create new process.") bulk_logger.info( "The environment variable PF_WORKER_COUNT is not set or invalid. Calculate the worker count based " - "on the currently memory usage") + "on the currently memory usage.") bulk_logger.info( f"Calculated process count ({worker_count}) by taking the minimum value among estimated process " f"count ({estimated_available_worker_count}), the row count ({self._nlines}) and the default " f"worker count ({self._DEFAULT_WORKER_COUNT}).") return worker_count else: - bulk_logger.info("Using fork to create new process") + bulk_logger.info("Using fork to create new process.") bulk_logger.info( f"Calculated process count ({worker_count}) by taking the minimum value among the " - f"default worker_count ({self._DEFAULT_WORKER_COUNT}) and the row count ({self._nlines})") + f"default worker_count ({self._DEFAULT_WORKER_COUNT}) and the row count ({self._nlines}).") return worker_count diff --git a/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py b/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py index 1a44d92b724..1541bd4b7bc 100644 --- a/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py +++ b/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py @@ -82,12 +82,12 @@ def test_fork_mode_parallelism_in_subprocess( assert pool._n_process == n_process if is_set_environ_pf_worker_count: mock_logger.info.assert_any_call( - f"Process count set to {pf_worker_count} based on 'PF_WORKER_COUNT' environment variable.") + f"Process count set to {pf_worker_count} based on the 'PF_WORKER_COUNT' environment variable.") else: - mock_logger.info.assert_any_call("Using fork to create new process") + mock_logger.info.assert_any_call("Using fork to create new process.") mock_logger.info.assert_any_call( f"Calculated process count ({n_process}) by taking the minimum value among the " - f"default worker_count ({pool._DEFAULT_WORKER_COUNT}) and the row count ({nlines})" + f"default worker_count ({pool._DEFAULT_WORKER_COUNT}) and the row count ({nlines})." ) @@ -128,19 +128,21 @@ def test_spawn_mode_parallelism_in_subprocess( assert pool._n_process == n_process if is_set_environ_pf_worker_count and is_calculation_smaller_than_set: mock_logger.info.assert_any_call( - f"Process count set to {pf_worker_count} based on 'PF_WORKER_COUNT' environment variable.") + f"Process count set to {pf_worker_count} based on the 'PF_WORKER_COUNT' " + f"environment variable.") mock_logger.warning.assert_any_call( f"The estimated available worker count calculated based on the system available memory " f"is {estimated_available_worker_count}, but the PF_WORKER_COUNT is set to " f"{pf_worker_count}. This may affect optimal memory usage and performance. ") elif is_set_environ_pf_worker_count and not is_calculation_smaller_than_set: mock_logger.info.assert_any_call( - f"Process count set to {pf_worker_count} based on 'PF_WORKER_COUNT' environment variable.") + f"Process count set to {pf_worker_count} based on the 'PF_WORKER_COUNT' " + f"environment variable.") elif not is_set_environ_pf_worker_count: - mock_logger.info.assert_any_call("Not using fork to create new process") + mock_logger.info.assert_any_call("Not using fork to create new process.") mock_logger.info.assert_any_call( "The environment variable PF_WORKER_COUNT is not set or invalid. Calculate the worker " - "count based on the currently memory usage" + "count based on the currently memory usage." ) mock_logger.info.assert_any_call( f"Calculated process count ({n_process}) by taking the minimum value among estimated " From 96216c996435edee02a975dda7a511ee447c23b6 Mon Sep 17 00:00:00 2001 From: LuLu Zuo Date: Wed, 6 Dec 2023 11:53:03 +0800 Subject: [PATCH 23/32] Refine _determine_worker_count --- .../executor/_line_execution_process_pool.py | 67 ++++++++++++------- .../test_line_execution_process_pool.py | 2 +- 2 files changed, 43 insertions(+), 26 deletions(-) diff --git a/src/promptflow/promptflow/executor/_line_execution_process_pool.py b/src/promptflow/promptflow/executor/_line_execution_process_pool.py index 39ff7c1dcaf..6dd900aae75 100644 --- a/src/promptflow/promptflow/executor/_line_execution_process_pool.py +++ b/src/promptflow/promptflow/executor/_line_execution_process_pool.py @@ -406,35 +406,52 @@ def _generate_thread_status_messages(self, pool: ThreadPool, total_count: int): def _determine_worker_count(self): worker_count = get_int_env_var("PF_WORKER_COUNT") + + # Starting a new process in non-fork mode requires to allocate memory. Calculate the maximum number of processes + # based on available memory to avoid memory bursting. estimated_available_worker_count = get_available_max_worker_count() if not self._use_fork else None + + # If the environment variable PF_WORKER_COUNT exists ant valid, use the value as the worker_count. if worker_count is not None and worker_count > 0: - bulk_logger.info( - f"Process count set to {worker_count} based on the 'PF_WORKER_COUNT' environment variable.") - if not self._use_fork and estimated_available_worker_count < worker_count: - bulk_logger.warning( - f"The estimated available worker count calculated based on the system available memory " - f"is {estimated_available_worker_count}, but the PF_WORKER_COUNT is set to " - f"{worker_count}. This may affect optimal memory usage and performance. ") + self._log_set_worker_count(worker_count, estimated_available_worker_count) return worker_count + + return self._calculate_worker_count(estimated_available_worker_count) + + def _log_set_worker_count(self, worker_count, estimated_available_worker_count): + bulk_logger.info( + f"Process count set to {worker_count} based on the 'PF_WORKER_COUNT' environment variable.") + if estimated_available_worker_count is not None and estimated_available_worker_count < worker_count: + bulk_logger.warning( + f"The estimated available worker count calculated based on the system available memory " + f"is {estimated_available_worker_count}, but the PF_WORKER_COUNT is set to " + f"{worker_count}. This may affect optimal memory usage and performance.") + + def _calculate_worker_count(self, estimated_available_worker_count): + if self._use_fork: + return self._calculate_worker_count_with_fork() else: - worker_count = min(self._DEFAULT_WORKER_COUNT, self._nlines) - if not self._use_fork: - worker_count = min(worker_count, estimated_available_worker_count) - bulk_logger.info("Not using fork to create new process.") - bulk_logger.info( - "The environment variable PF_WORKER_COUNT is not set or invalid. Calculate the worker count based " - "on the currently memory usage.") - bulk_logger.info( - f"Calculated process count ({worker_count}) by taking the minimum value among estimated process " - f"count ({estimated_available_worker_count}), the row count ({self._nlines}) and the default " - f"worker count ({self._DEFAULT_WORKER_COUNT}).") - return worker_count - else: - bulk_logger.info("Using fork to create new process.") - bulk_logger.info( - f"Calculated process count ({worker_count}) by taking the minimum value among the " - f"default worker_count ({self._DEFAULT_WORKER_COUNT}) and the row count ({self._nlines}).") - return worker_count + return self._calculate_worker_count_without_fork(estimated_available_worker_count) + + def _calculate_worker_count_with_fork(self): + worker_count = min(self._DEFAULT_WORKER_COUNT, self._nlines) + bulk_logger.info("Using fork to create new process.") + bulk_logger.info( + f"Calculated process count ({worker_count}) by taking the minimum value among the " + f"default worker_count ({self._DEFAULT_WORKER_COUNT}) and the row count ({self._nlines}).") + return worker_count + + def _calculate_worker_count_without_fork(self, estimated_available_worker_count): + worker_count = min(self._DEFAULT_WORKER_COUNT, self._nlines, estimated_available_worker_count) + bulk_logger.info("Not using fork to create new process.") + bulk_logger.info( + "The environment variable PF_WORKER_COUNT is not set or invalid. Calculate the worker count based " + "on the currently memory usage.") + bulk_logger.info( + f"Calculated process count ({worker_count}) by taking the minimum value among estimated process " + f"count ({estimated_available_worker_count}), the row count ({self._nlines}) and the default " + f"worker count ({self._DEFAULT_WORKER_COUNT}).") + return worker_count def _exec_line( diff --git a/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py b/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py index 1541bd4b7bc..584ec31ff3c 100644 --- a/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py +++ b/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py @@ -133,7 +133,7 @@ def test_spawn_mode_parallelism_in_subprocess( mock_logger.warning.assert_any_call( f"The estimated available worker count calculated based on the system available memory " f"is {estimated_available_worker_count}, but the PF_WORKER_COUNT is set to " - f"{pf_worker_count}. This may affect optimal memory usage and performance. ") + f"{pf_worker_count}. This may affect optimal memory usage and performance.") elif is_set_environ_pf_worker_count and not is_calculation_smaller_than_set: mock_logger.info.assert_any_call( f"Process count set to {pf_worker_count} based on the 'PF_WORKER_COUNT' " From 8882ba2cd9343c187722a7deb024e64d917ba419 Mon Sep 17 00:00:00 2001 From: LuLu Zuo Date: Wed, 6 Dec 2023 14:12:51 +0800 Subject: [PATCH 24/32] fix typo --- .../promptflow/executor/_line_execution_process_pool.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/promptflow/promptflow/executor/_line_execution_process_pool.py b/src/promptflow/promptflow/executor/_line_execution_process_pool.py index 6dd900aae75..6fc5dd8b9eb 100644 --- a/src/promptflow/promptflow/executor/_line_execution_process_pool.py +++ b/src/promptflow/promptflow/executor/_line_execution_process_pool.py @@ -411,7 +411,7 @@ def _determine_worker_count(self): # based on available memory to avoid memory bursting. estimated_available_worker_count = get_available_max_worker_count() if not self._use_fork else None - # If the environment variable PF_WORKER_COUNT exists ant valid, use the value as the worker_count. + # If the environment variable PF_WORKER_COUNT exists and valid, use the value as the worker_count. if worker_count is not None and worker_count > 0: self._log_set_worker_count(worker_count, estimated_available_worker_count) return worker_count From 05ed6c19a481c953862b62a64fc4b0fa5df3c59c Mon Sep 17 00:00:00 2001 From: LuLu Zuo Date: Wed, 6 Dec 2023 16:45:29 +0800 Subject: [PATCH 25/32] Optimization logs --- .../executor/_line_execution_process_pool.py | 23 ++++++++----------- .../test_line_execution_process_pool.py | 21 +++++++++-------- 2 files changed, 21 insertions(+), 23 deletions(-) diff --git a/src/promptflow/promptflow/executor/_line_execution_process_pool.py b/src/promptflow/promptflow/executor/_line_execution_process_pool.py index 6fc5dd8b9eb..36977936466 100644 --- a/src/promptflow/promptflow/executor/_line_execution_process_pool.py +++ b/src/promptflow/promptflow/executor/_line_execution_process_pool.py @@ -416,28 +416,25 @@ def _determine_worker_count(self): self._log_set_worker_count(worker_count, estimated_available_worker_count) return worker_count - return self._calculate_worker_count(estimated_available_worker_count) + if not self._use_fork: + return self._calculate_worker_count_without_fork(estimated_available_worker_count) + else: + return self._calculate_worker_count_with_fork() def _log_set_worker_count(self, worker_count, estimated_available_worker_count): bulk_logger.info( - f"Process count set to {worker_count} based on the 'PF_WORKER_COUNT' environment variable.") + f"Set process count to {worker_count} with the environment variable 'PF_WORKER_COUNT'.") if estimated_available_worker_count is not None and estimated_available_worker_count < worker_count: bulk_logger.warning( - f"The estimated available worker count calculated based on the system available memory " - f"is {estimated_available_worker_count}, but the PF_WORKER_COUNT is set to " - f"{worker_count}. This may affect optimal memory usage and performance.") - - def _calculate_worker_count(self, estimated_available_worker_count): - if self._use_fork: - return self._calculate_worker_count_with_fork() - else: - return self._calculate_worker_count_without_fork(estimated_available_worker_count) + f"The current process count({worker_count}) is larger than recommended process count" + f"({estimated_available_worker_count}) that estimated by system available memory. This may " + f"cause memory exhaustion") def _calculate_worker_count_with_fork(self): worker_count = min(self._DEFAULT_WORKER_COUNT, self._nlines) bulk_logger.info("Using fork to create new process.") bulk_logger.info( - f"Calculated process count ({worker_count}) by taking the minimum value among the " + f"Set process count to ({worker_count}) by taking the minimum value among the " f"default worker_count ({self._DEFAULT_WORKER_COUNT}) and the row count ({self._nlines}).") return worker_count @@ -448,7 +445,7 @@ def _calculate_worker_count_without_fork(self, estimated_available_worker_count) "The environment variable PF_WORKER_COUNT is not set or invalid. Calculate the worker count based " "on the currently memory usage.") bulk_logger.info( - f"Calculated process count ({worker_count}) by taking the minimum value among estimated process " + f"Set process count to ({worker_count}) by taking the minimum value among estimated process " f"count ({estimated_available_worker_count}), the row count ({self._nlines}) and the default " f"worker count ({self._DEFAULT_WORKER_COUNT}).") return worker_count diff --git a/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py b/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py index 584ec31ff3c..9452f60ed29 100644 --- a/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py +++ b/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py @@ -82,11 +82,12 @@ def test_fork_mode_parallelism_in_subprocess( assert pool._n_process == n_process if is_set_environ_pf_worker_count: mock_logger.info.assert_any_call( - f"Process count set to {pf_worker_count} based on the 'PF_WORKER_COUNT' environment variable.") + f"Set process count to {pf_worker_count} with the environment " + f"variable 'PF_WORKER_COUNT'.") else: mock_logger.info.assert_any_call("Using fork to create new process.") mock_logger.info.assert_any_call( - f"Calculated process count ({n_process}) by taking the minimum value among the " + f"Set process count to ({n_process}) by taking the minimum value among the " f"default worker_count ({pool._DEFAULT_WORKER_COUNT}) and the row count ({nlines})." ) @@ -128,16 +129,16 @@ def test_spawn_mode_parallelism_in_subprocess( assert pool._n_process == n_process if is_set_environ_pf_worker_count and is_calculation_smaller_than_set: mock_logger.info.assert_any_call( - f"Process count set to {pf_worker_count} based on the 'PF_WORKER_COUNT' " - f"environment variable.") + f"Set process count to {pf_worker_count} with the environment " + f"variable 'PF_WORKER_COUNT'.") mock_logger.warning.assert_any_call( - f"The estimated available worker count calculated based on the system available memory " - f"is {estimated_available_worker_count}, but the PF_WORKER_COUNT is set to " - f"{pf_worker_count}. This may affect optimal memory usage and performance.") + f"The current process count({pf_worker_count}) is larger than recommended process count" + f"({estimated_available_worker_count}) that estimated by system available memory. This may " + f"cause memory exhaustion") elif is_set_environ_pf_worker_count and not is_calculation_smaller_than_set: mock_logger.info.assert_any_call( - f"Process count set to {pf_worker_count} based on the 'PF_WORKER_COUNT' " - f"environment variable.") + f"Set process count to {pf_worker_count} with the environment " + f"variable 'PF_WORKER_COUNT'.") elif not is_set_environ_pf_worker_count: mock_logger.info.assert_any_call("Not using fork to create new process.") mock_logger.info.assert_any_call( @@ -145,7 +146,7 @@ def test_spawn_mode_parallelism_in_subprocess( "count based on the currently memory usage." ) mock_logger.info.assert_any_call( - f"Calculated process count ({n_process}) by taking the minimum value among estimated " + f"Set process count to ({n_process}) by taking the minimum value among estimated " f"process count ({estimated_available_worker_count}), the row count ({nlines}) and the " f"default worker count ({pool._DEFAULT_WORKER_COUNT})." ) From 01653a83e61c862b766d1282d695c943f66211df Mon Sep 17 00:00:00 2001 From: LuLu Zuo Date: Wed, 6 Dec 2023 16:55:35 +0800 Subject: [PATCH 26/32] Rmove bracket --- .../promptflow/executor/_line_execution_process_pool.py | 4 ++-- .../unittests/processpool/test_line_execution_process_pool.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/promptflow/promptflow/executor/_line_execution_process_pool.py b/src/promptflow/promptflow/executor/_line_execution_process_pool.py index 36977936466..620c9e86598 100644 --- a/src/promptflow/promptflow/executor/_line_execution_process_pool.py +++ b/src/promptflow/promptflow/executor/_line_execution_process_pool.py @@ -434,7 +434,7 @@ def _calculate_worker_count_with_fork(self): worker_count = min(self._DEFAULT_WORKER_COUNT, self._nlines) bulk_logger.info("Using fork to create new process.") bulk_logger.info( - f"Set process count to ({worker_count}) by taking the minimum value among the " + f"Set process count to {worker_count} by taking the minimum value among the " f"default worker_count ({self._DEFAULT_WORKER_COUNT}) and the row count ({self._nlines}).") return worker_count @@ -445,7 +445,7 @@ def _calculate_worker_count_without_fork(self, estimated_available_worker_count) "The environment variable PF_WORKER_COUNT is not set or invalid. Calculate the worker count based " "on the currently memory usage.") bulk_logger.info( - f"Set process count to ({worker_count}) by taking the minimum value among estimated process " + f"Set process count to {worker_count} by taking the minimum value among estimated process " f"count ({estimated_available_worker_count}), the row count ({self._nlines}) and the default " f"worker count ({self._DEFAULT_WORKER_COUNT}).") return worker_count diff --git a/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py b/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py index 9452f60ed29..185446d1f63 100644 --- a/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py +++ b/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py @@ -87,7 +87,7 @@ def test_fork_mode_parallelism_in_subprocess( else: mock_logger.info.assert_any_call("Using fork to create new process.") mock_logger.info.assert_any_call( - f"Set process count to ({n_process}) by taking the minimum value among the " + f"Set process count to {n_process} by taking the minimum value among the " f"default worker_count ({pool._DEFAULT_WORKER_COUNT}) and the row count ({nlines})." ) @@ -146,7 +146,7 @@ def test_spawn_mode_parallelism_in_subprocess( "count based on the currently memory usage." ) mock_logger.info.assert_any_call( - f"Set process count to ({n_process}) by taking the minimum value among estimated " + f"Set process count to {n_process} by taking the minimum value among estimated " f"process count ({estimated_available_worker_count}), the row count ({nlines}) and the " f"default worker count ({pool._DEFAULT_WORKER_COUNT})." ) From 130b1b166ca752d2cd6936fcdc89d2e5d4523621 Mon Sep 17 00:00:00 2001 From: LuLu Zuo Date: Wed, 6 Dec 2023 17:00:54 +0800 Subject: [PATCH 27/32] Add space --- .../promptflow/executor/_line_execution_process_pool.py | 2 +- .../unittests/processpool/test_line_execution_process_pool.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/promptflow/promptflow/executor/_line_execution_process_pool.py b/src/promptflow/promptflow/executor/_line_execution_process_pool.py index 620c9e86598..3be9c8e2421 100644 --- a/src/promptflow/promptflow/executor/_line_execution_process_pool.py +++ b/src/promptflow/promptflow/executor/_line_execution_process_pool.py @@ -426,7 +426,7 @@ def _log_set_worker_count(self, worker_count, estimated_available_worker_count): f"Set process count to {worker_count} with the environment variable 'PF_WORKER_COUNT'.") if estimated_available_worker_count is not None and estimated_available_worker_count < worker_count: bulk_logger.warning( - f"The current process count({worker_count}) is larger than recommended process count" + f"The current process count ({worker_count}) is larger than recommended process count " f"({estimated_available_worker_count}) that estimated by system available memory. This may " f"cause memory exhaustion") diff --git a/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py b/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py index 185446d1f63..3230b32e7df 100644 --- a/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py +++ b/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py @@ -132,7 +132,7 @@ def test_spawn_mode_parallelism_in_subprocess( f"Set process count to {pf_worker_count} with the environment " f"variable 'PF_WORKER_COUNT'.") mock_logger.warning.assert_any_call( - f"The current process count({pf_worker_count}) is larger than recommended process count" + f"The current process count ({pf_worker_count}) is larger than recommended process count " f"({estimated_available_worker_count}) that estimated by system available memory. This may " f"cause memory exhaustion") elif is_set_environ_pf_worker_count and not is_calculation_smaller_than_set: From beff59ddd1ddb3f55e2a6ab6301c9167b2fca9a9 Mon Sep 17 00:00:00 2001 From: LuLu Zuo Date: Wed, 6 Dec 2023 21:21:28 +0800 Subject: [PATCH 28/32] Optimization log --- .../executor/_line_execution_process_pool.py | 16 +++++++++------- .../test_line_execution_process_pool.py | 19 ++++++++++--------- 2 files changed, 19 insertions(+), 16 deletions(-) diff --git a/src/promptflow/promptflow/executor/_line_execution_process_pool.py b/src/promptflow/promptflow/executor/_line_execution_process_pool.py index 3be9c8e2421..0d6405a91b5 100644 --- a/src/promptflow/promptflow/executor/_line_execution_process_pool.py +++ b/src/promptflow/promptflow/executor/_line_execution_process_pool.py @@ -589,14 +589,16 @@ def get_available_max_worker_count(): # 2. When the degree of parallelism is 1, main process executes the task directly and not # create the child process bulk_logger.warning( - f"Available max worker count {estimated_available_worker_count} is less than 1, set it to 1.") + f"Current system's available memory is {available_memory}MB, less than the memory " + f"{process_memory}MB required by the process. The maximum available worker count is 1.") estimated_available_worker_count = 1 - bulk_logger.info( - f"Current system's available memory is {available_memory}MB, " - f"memory consumption of current process is {process_memory}MB, " - f"estimated available worker count is {available_memory}/{process_memory} " - f"= {estimated_available_worker_count}" - ) + else: + bulk_logger.info( + f"Current system's available memory is {available_memory}MB, " + f"memory consumption of current process is {process_memory}MB, " + f"estimated available worker count is {available_memory}/{process_memory} " + f"= {estimated_available_worker_count}" + ) return estimated_available_worker_count diff --git a/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py b/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py index 3230b32e7df..37626d02514 100644 --- a/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py +++ b/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py @@ -436,7 +436,7 @@ class TestGetAvailableMaxWorkerCount: "available_memory, process_memory, expected_max_worker_count, actual_calculate_worker_count", [ (128.0, 64.0, 2, 2), # available_memory/process_memory > 1 - (63.0, 64.0, 1, 1), # available_memory/process_memory < 1 + (63.0, 64.0, 1, 0), # available_memory/process_memory < 1 ], ) def test_get_available_max_worker_count( @@ -452,12 +452,13 @@ def test_get_available_max_worker_count( assert estimated_available_worker_count == expected_max_worker_count if actual_calculate_worker_count < 1: mock_logger.warning.assert_called_with( - f"Available max worker count {actual_calculate_worker_count} is less than 1, " - "set it to 1." + f"Current system's available memory is {available_memory}MB, less than the memory " + f"{process_memory}MB required by the process. The maximum available worker count is 1." + ) + else: + mock_logger.info.assert_called_with( + f"Current system's available memory is {available_memory}MB, " + f"memory consumption of current process is {process_memory}MB, " + f"estimated available worker count is {available_memory}/{process_memory} " + f"= {actual_calculate_worker_count}" ) - mock_logger.info.assert_called_with( - f"Current system's available memory is {available_memory}MB, " - f"memory consumption of current process is {process_memory}MB, " - f"estimated available worker count is {available_memory}/{process_memory} " - f"= {actual_calculate_worker_count}" - ) From 1ac33214093a46cf42342d91ccc480fa808f9b20 Mon Sep 17 00:00:00 2001 From: LuLu Zuo Date: Thu, 7 Dec 2023 11:29:19 +0800 Subject: [PATCH 29/32] Refine _determine_worker_count logic --- .../executor/_line_execution_process_pool.py | 49 +++++++++---------- .../test_line_execution_process_pool.py | 22 +++++---- 2 files changed, 34 insertions(+), 37 deletions(-) diff --git a/src/promptflow/promptflow/executor/_line_execution_process_pool.py b/src/promptflow/promptflow/executor/_line_execution_process_pool.py index 0d6405a91b5..6e3bd3b083b 100644 --- a/src/promptflow/promptflow/executor/_line_execution_process_pool.py +++ b/src/promptflow/promptflow/executor/_line_execution_process_pool.py @@ -406,20 +406,35 @@ def _generate_thread_status_messages(self, pool: ThreadPool, total_count: int): def _determine_worker_count(self): worker_count = get_int_env_var("PF_WORKER_COUNT") + estimated_available_worker_count = None - # Starting a new process in non-fork mode requires to allocate memory. Calculate the maximum number of processes - # based on available memory to avoid memory bursting. - estimated_available_worker_count = get_available_max_worker_count() if not self._use_fork else None + if not self._use_fork: + # Starting a new process in non-fork mode requires to allocate memory. Calculate the maximum number of + # processes based on available memory to avoid memory bursting. + estimated_available_worker_count = get_available_max_worker_count() # If the environment variable PF_WORKER_COUNT exists and valid, use the value as the worker_count. if worker_count is not None and worker_count > 0: self._log_set_worker_count(worker_count, estimated_available_worker_count) return worker_count - if not self._use_fork: - return self._calculate_worker_count_without_fork(estimated_available_worker_count) - else: - return self._calculate_worker_count_with_fork() + # If the environment variable PF_WORKER_COUNT is not set or invalid, take the minimum value among the + # factors: default_worker_count, row_count and estimated_worker_count_based_on_memory_usage + factors = { + "default_worker_count": self._DEFAULT_WORKER_COUNT, + "row_count": self._nlines, + } + + if estimated_available_worker_count is not None and estimated_available_worker_count > 0: + factors.update({ + "estimated_worker_count_based_on_memory_usage": estimated_available_worker_count, + }) + + # Take the minimum value as the result + worker_count = min(factors.values()) + bulk_logger.info( + f"Set process count to {worker_count} by taking the minimum value among the factors of {factors}.") + return worker_count def _log_set_worker_count(self, worker_count, estimated_available_worker_count): bulk_logger.info( @@ -430,26 +445,6 @@ def _log_set_worker_count(self, worker_count, estimated_available_worker_count): f"({estimated_available_worker_count}) that estimated by system available memory. This may " f"cause memory exhaustion") - def _calculate_worker_count_with_fork(self): - worker_count = min(self._DEFAULT_WORKER_COUNT, self._nlines) - bulk_logger.info("Using fork to create new process.") - bulk_logger.info( - f"Set process count to {worker_count} by taking the minimum value among the " - f"default worker_count ({self._DEFAULT_WORKER_COUNT}) and the row count ({self._nlines}).") - return worker_count - - def _calculate_worker_count_without_fork(self, estimated_available_worker_count): - worker_count = min(self._DEFAULT_WORKER_COUNT, self._nlines, estimated_available_worker_count) - bulk_logger.info("Not using fork to create new process.") - bulk_logger.info( - "The environment variable PF_WORKER_COUNT is not set or invalid. Calculate the worker count based " - "on the currently memory usage.") - bulk_logger.info( - f"Set process count to {worker_count} by taking the minimum value among estimated process " - f"count ({estimated_available_worker_count}), the row count ({self._nlines}) and the default " - f"worker count ({self._DEFAULT_WORKER_COUNT}).") - return worker_count - def _exec_line( executor: FlowExecutor, diff --git a/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py b/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py index 37626d02514..313e130b953 100644 --- a/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py +++ b/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py @@ -85,10 +85,13 @@ def test_fork_mode_parallelism_in_subprocess( f"Set process count to {pf_worker_count} with the environment " f"variable 'PF_WORKER_COUNT'.") else: - mock_logger.info.assert_any_call("Using fork to create new process.") + factors = { + "default_worker_count": pool._DEFAULT_WORKER_COUNT, + "row_count": pool._nlines, + } mock_logger.info.assert_any_call( f"Set process count to {n_process} by taking the minimum value among the " - f"default worker_count ({pool._DEFAULT_WORKER_COUNT}) and the row count ({nlines})." + f"factors of {factors}." ) @@ -126,6 +129,7 @@ def test_spawn_mode_parallelism_in_subprocess( False, None, ) as pool: + assert pool._n_process == n_process if is_set_environ_pf_worker_count and is_calculation_smaller_than_set: mock_logger.info.assert_any_call( @@ -140,15 +144,13 @@ def test_spawn_mode_parallelism_in_subprocess( f"Set process count to {pf_worker_count} with the environment " f"variable 'PF_WORKER_COUNT'.") elif not is_set_environ_pf_worker_count: - mock_logger.info.assert_any_call("Not using fork to create new process.") - mock_logger.info.assert_any_call( - "The environment variable PF_WORKER_COUNT is not set or invalid. Calculate the worker " - "count based on the currently memory usage." - ) + factors = { + "default_worker_count": pool._DEFAULT_WORKER_COUNT, + "row_count": pool._nlines, + "estimated_worker_count_based_on_memory_usage": estimated_available_worker_count + } mock_logger.info.assert_any_call( - f"Set process count to {n_process} by taking the minimum value among estimated " - f"process count ({estimated_available_worker_count}), the row count ({nlines}) and the " - f"default worker count ({pool._DEFAULT_WORKER_COUNT})." + f"Set process count to {n_process} by taking the minimum value among the factors of {factors}." ) From 3c3cb174cf2b8300a5ca40331d5135187933b811 Mon Sep 17 00:00:00 2001 From: LuLu Zuo Date: Thu, 7 Dec 2023 11:31:30 +0800 Subject: [PATCH 30/32] fix flask8 --- .../unittests/processpool/test_line_execution_process_pool.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py b/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py index 313e130b953..05607ed9b51 100644 --- a/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py +++ b/src/promptflow/tests/executor/unittests/processpool/test_line_execution_process_pool.py @@ -150,7 +150,8 @@ def test_spawn_mode_parallelism_in_subprocess( "estimated_worker_count_based_on_memory_usage": estimated_available_worker_count } mock_logger.info.assert_any_call( - f"Set process count to {n_process} by taking the minimum value among the factors of {factors}." + f"Set process count to {n_process} by taking the minimum value among the factors " + f"of {factors}." ) From 2c1a1b6dc6810b69b6955578c4dd017e8c90f43a Mon Sep 17 00:00:00 2001 From: LuLu Zuo Date: Thu, 7 Dec 2023 12:00:07 +0800 Subject: [PATCH 31/32] Refine if logic --- .../promptflow/executor/_line_execution_process_pool.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/src/promptflow/promptflow/executor/_line_execution_process_pool.py b/src/promptflow/promptflow/executor/_line_execution_process_pool.py index 6e3bd3b083b..e7bc3fc140e 100644 --- a/src/promptflow/promptflow/executor/_line_execution_process_pool.py +++ b/src/promptflow/promptflow/executor/_line_execution_process_pool.py @@ -406,12 +406,10 @@ def _generate_thread_status_messages(self, pool: ThreadPool, total_count: int): def _determine_worker_count(self): worker_count = get_int_env_var("PF_WORKER_COUNT") - estimated_available_worker_count = None - if not self._use_fork: - # Starting a new process in non-fork mode requires to allocate memory. Calculate the maximum number of - # processes based on available memory to avoid memory bursting. - estimated_available_worker_count = get_available_max_worker_count() + # Starting a new process in non-fork mode requires to allocate memory. Calculate the maximum number of processes + # based on available memory to avoid memory bursting. + estimated_available_worker_count = get_available_max_worker_count() if not self._use_fork else None # If the environment variable PF_WORKER_COUNT exists and valid, use the value as the worker_count. if worker_count is not None and worker_count > 0: From c1e228a1244a4011d0c022c0b3ec333e7944b581 Mon Sep 17 00:00:00 2001 From: LuLu Zuo Date: Thu, 7 Dec 2023 12:59:25 +0800 Subject: [PATCH 32/32] Redefine the calculation of factors --- .../executor/_line_execution_process_pool.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/src/promptflow/promptflow/executor/_line_execution_process_pool.py b/src/promptflow/promptflow/executor/_line_execution_process_pool.py index e7bc3fc140e..3e8f50f5e68 100644 --- a/src/promptflow/promptflow/executor/_line_execution_process_pool.py +++ b/src/promptflow/promptflow/executor/_line_execution_process_pool.py @@ -421,17 +421,15 @@ def _determine_worker_count(self): factors = { "default_worker_count": self._DEFAULT_WORKER_COUNT, "row_count": self._nlines, + "estimated_worker_count_based_on_memory_usage": estimated_available_worker_count, } - if estimated_available_worker_count is not None and estimated_available_worker_count > 0: - factors.update({ - "estimated_worker_count_based_on_memory_usage": estimated_available_worker_count, - }) + valid_factors = {k: v for k, v in factors.items() if v is not None and v > 0} # Take the minimum value as the result - worker_count = min(factors.values()) + worker_count = min(valid_factors.values()) bulk_logger.info( - f"Set process count to {worker_count} by taking the minimum value among the factors of {factors}.") + f"Set process count to {worker_count} by taking the minimum value among the factors of {valid_factors}.") return worker_count def _log_set_worker_count(self, worker_count, estimated_available_worker_count):