Repository: microsoft/promptflow
File: src/promptflow-core/promptflow/executor/_service/utils/process_utils.py (Python)

Code before:
                    raise UnexpectedError(
                        message="Unexpected error occurred while executing the request",
                        target=ErrorTarget.EXECUTOR,
                    )
                # JsonSerializedPromptflowException will be raised here
                # no need to change to PromptflowException since it will be handled in app.exception_handler
                raise exception

            service_logger.info(f"[{p.pid}--{os.getpid()}] Process finished.")
            return return_dict.get("result", {})
        finally:
            if run_id:
                ProcessManager().remove_process(run_id)


def _set_environment_variables_then_execute_target_function(
    target_function: Callable,
    args: tuple,
    kwargs: dict,
    return_dict: dict,
    error_dict: dict,
    context_dict: dict,
    environment_variables: dict,
):
    block_terminate_signal_to_parent()
    if environment_variables:
        service_logger.info("Set environment variable before invoke target function.")
        os.environ.update(environment_variables)
    with exception_wrapper(error_dict):
        if context_dict:
            OperationContext.get_instance().update(context_dict)
        service_logger.info("Start processing request in executor service...")
        result = target_function(*args, **kwargs)
        return_dict["result"] = result


@contextlib.contextmanager
def exception_wrapper(error_dict: dict):
    """Wrap the exception to a generic exception to avoid the pickle error."""
    try:
        yield
    except Exception as e:
        # Func runs in a child process, any customized exception can't have extra arguments other than message
        # Wrap the exception to a generic exception to avoid the pickle error
        # Ref: https://bugs.python.org/issue32696
        exception_dict = ExceptionPresenter.create(e).to_dict(include_debug_info=True)
        message = json.dumps(exception_dict)
        exception = JsonSerializedPromptflowException(message=message)
        error_dict["error"] = exception
        raise exception from e


Review comment (suggestion):
There is a util function `set_environment_variables` in `src\promptflow-core\promptflow\executor\_service\utils\service_utils.py`. You can use it.

Code after:
                        raise ExecutionCanceledError(run_id)
                    raise UnexpectedError(
                        message="Unexpected error occurred while executing the request",
                        target=ErrorTarget.EXECUTOR,
                    )
                # JsonSerializedPromptflowException will be raised here
                # no need to change to PromptflowException since it will be handled in app.exception_handler
                raise exception

            service_logger.info(f"[{p.pid}--{os.getpid()}] Process finished.")
            return return_dict.get("result", {})
        finally:
            if run_id:
                ProcessManager().remove_process(run_id)


def _execute_target_function(
    target_function: Callable,
    args: tuple,
    kwargs: dict,
    return_dict: dict,
    error_dict: dict,
    context_dict: dict,
    environment_variables: Mapping[str, Any],
):
    block_terminate_signal_to_parent()
    set_environment_variables(environment_variables)
    with exception_wrapper(error_dict):
        if context_dict:
            OperationContext.get_instance().update(context_dict)
        service_logger.info("Start processing request in executor service...")
        result = target_function(*args, **kwargs)
        return_dict["result"] = result


@contextlib.contextmanager
def exception_wrapper(error_dict: dict):
    """Wrap the exception to a generic exception to avoid the pickle error."""
    try:
        yield
    except Exception as e:
        # Func runs in a child process, any customized exception can't have extra arguments other than message
        # Wrap the exception to a generic exception to avoid the pickle error
        # Ref: https://bugs.python.org/issue32696
        exception_dict = ExceptionPresenter.create(e).to_dict(include_debug_info=True)
        message = json.dumps(exception_dict)
        exception = JsonSerializedPromptflowException(message=message)
        error_dict["error"] = exception
        raise exception from e