""" TaskVineExecutor utilizes the TaskVine distributed framework developed by the Cooperative Computing Lab (CCL) at Notre Dame to provide a fault-tolerant, high-throughput system for delegating Parsl tasks to thousands of remote machines """ # Import Python built-in libraries import threading import multiprocessing import logging import tempfile import hashlib import subprocess import os import queue import inspect import shutil import itertools import uuid from concurrent.futures import Future from typing import List, Optional, Union, Literal # Import Parsl constructs import parsl.utils as putils from parsl.data_provider.staging import Staging from parsl.serialize import serialize, deserialize from parsl.data_provider.files import File from parsl.errors import OptionalModuleMissing from parsl.providers.base import ExecutionProvider from parsl.providers import LocalProvider, CondorProvider from parsl.process_loggers import wrap_with_logs from parsl.addresses import get_any_address from parsl.executors.errors import ExecutorError from parsl.executors.status_handling import BlockProviderExecutor from parsl.executors.taskvine import exec_parsl_function from parsl.executors.taskvine.manager_config import TaskVineManagerConfig from parsl.executors.taskvine.factory_config import TaskVineFactoryConfig from parsl.executors.taskvine.errors import TaskVineTaskFailure from parsl.executors.taskvine.errors import TaskVineManagerFailure from parsl.executors.taskvine.utils import ParslTaskToVine from parsl.executors.taskvine.utils import ParslFileToVine from parsl.executors.taskvine.manager import _taskvine_submit_wait from parsl.executors.taskvine.factory import _taskvine_factory # Import other libraries import typeguard logger = logging.getLogger(__name__) class TaskVineExecutor(BlockProviderExecutor, putils.RepresentationMixin): """Executor to use TaskVine dynamic workflow system The TaskVineExecutor system utilizes the TaskVine framework to efficiently delegate Parsl apps to remote machines in clusters and grids using a fault-tolerant system. Users can run the vine_worker program on remote machines to connect to the TaskVineExecutor, and Parsl apps will then be sent out to these machines for execution and retrieval. This Executor sets up configurations for the TaskVine manager, TaskVine factory, and run both in separate processes. Sending tasks and receiving results are done through multiprocessing module native to Python. Parameters ---------- label: str A human readable label for the executor, unique with respect to other executors. Default is "TaskVineExecutor". worker_launch_method: Union[Literal['provider'], Literal['factory'], Literal['manual']] Choose to use Parsl provider, TaskVine factory, or manual user-provided workers to scale workers. Options are among {'provider', 'factory', 'manual'}. Default is 'factory'. function_exec_mode: Union[Literal['regular'], Literal['serverless']] Choose to execute functions with a regular fresh python process or a pre-warmed forked python process. Default is 'regular'. manager_config: TaskVineManagerConfig Configuration for the TaskVine manager. Default factory_config: TaskVineFactoryConfig Configuration for the TaskVine factory. Use of factory is disabled by default. provider: ExecutionProvider The Parsl provider that will spawn worker processes. Default to spawning one local vine worker process. storage_access: List[Staging] Define Parsl file staging providers for this executor. Default is None. """ radio_mode = "filesystem" @typeguard.typechecked def __init__(self, label: str = "TaskVineExecutor", worker_launch_method: Union[Literal['provider'], Literal['factory'], Literal['manual']] = 'factory', function_exec_mode: Union[Literal['regular'], Literal['serverless']] = 'regular', manager_config: TaskVineManagerConfig = TaskVineManagerConfig(), factory_config: TaskVineFactoryConfig = TaskVineFactoryConfig(), provider: Optional[ExecutionProvider] = LocalProvider(init_blocks=1), storage_access: Optional[List[Staging]] = None): # Set worker launch option for this executor if worker_launch_method == 'factory' or worker_launch_method == 'manual': provider = None # Initialize the parent class with the execution provider and block error handling enabled. # If provider is None, then no worker is launched via the provider method. BlockProviderExecutor.__init__(self, provider=provider, block_error_handler=True) # Raise an exception if there's a problem importing TaskVine try: import ndcctools.taskvine logger.debug(f'TaskVine default port: {ndcctools.taskvine.cvine.VINE_DEFAULT_PORT}') except ImportError: raise OptionalModuleMissing(['taskvine'], "TaskVineExecutor requires the taskvine module.") # Executor configurations self.label = label self.worker_launch_method = worker_launch_method self.function_exec_mode = function_exec_mode self.manager_config = manager_config self.factory_config = factory_config self.storage_access = storage_access # Queue to send ready tasks from TaskVine executor process to TaskVine manager process self._ready_task_queue: multiprocessing.Queue = multiprocessing.Queue() # Queue to send finished tasks from TaskVine manager process to TaskVine executor process self._finished_task_queue: multiprocessing.Queue = multiprocessing.Queue() # Event to signal whether the manager and factory processes should stop running self._should_stop = multiprocessing.Event() # TaskVine manager process self._submit_process = None # TaskVine factory process self._factory_process = None # Executor thread to collect results from TaskVine manager and set # tasks' futures to done status. self._collector_thread = None # track task id of submitted parsl tasks # task ids are incremental and start from 0 self._executor_task_counter = 0 # track number of tasks that are waiting/running self._outstanding_tasks = 0 # Lock for threads to access self._outstanding_tasks attribute self._outstanding_tasks_lock = threading.Lock() # Threading lock to manage self.tasks dictionary object, which maps a task id # to its future object. self._tasks_lock = threading.Lock() # Worker command to be given to an execution provider (e.g., local or Condor) self._worker_command = "" # Path to directory that holds all tasks' data and results. self._function_data_dir = "" # helper scripts to prepare package tarballs for Parsl apps self._package_analyze_script = shutil.which("poncho_package_analyze") self._package_create_script = shutil.which("poncho_package_create") if self._package_analyze_script is None or self._package_create_script is None: self._poncho_available = False else: self._poncho_available = True def _get_launch_command(self, block_id): # Implements BlockProviderExecutor's abstract method. # This executor uses different terminology for worker/launch # commands than in htex. return f"PARSL_WORKER_BLOCK_ID={block_id} {self._worker_command}" def __synchronize_manager_factory_comm_settings(self): # Synchronize the communication settings between the manager and the factory # so the factory can direct workers to contact the manager. # If the manager can choose any available port (port number = 0), # then it must have a project name # so the factory can look it up. Otherwise the factory process will not know the # port number as it's only chosen when the TaskVine manager process is run. if self.manager_config.port == 0 and self.manager_config.project_name is None: self.manager_config.project_name = "parsl-vine-" + str(uuid.uuid4()) # guess the host name if the project name is not given if not self.manager_config.project_name: self.manager_config.address = get_any_address() # Factory communication settings are overridden by manager communication settings. self.factory_config._project_port = self.manager_config.port self.factory_config._project_address = self.manager_config.address self.factory_config._project_name = self.manager_config.project_name self.factory_config._project_password_file = self.manager_config.project_password_file logger.debug('Communication settings between TaskVine manager and factory synchronized.') def __create_data_and_logging_dirs(self): # Create neccessary data and logging directories # Use the current run directory from Parsl run_dir = self.run_dir # Create directories for data and results log_dir = os.path.join(run_dir, self.label) self._function_data_dir = os.path.join(run_dir, self.label, "function_data") os.makedirs(log_dir) os.makedirs(self._function_data_dir) # put TaskVine logs outside of a Parsl run as TaskVine caches between runs while # Parsl does not. vine_log_dir = os.path.join(os.path.dirname(run_dir), self.label) if self.manager_config.vine_log_dir is None: self.manager_config.vine_log_dir = vine_log_dir # factory logs go with manager logs regardless self.factory_config.scratch_dir = self.manager_config.vine_log_dir logger.debug(f"Function data directory: {self._function_data_dir}, log directory: {log_dir}") logger.debug(f"TaskVine manager log directory: {self.manager_config.vine_log_dir}, factory log directory: {self.factory_config.scratch_dir}") def start(self): """Create submit process and collector thread to create, send, and retrieve Parsl tasks within the TaskVine system. """ # Synchronize connection and communication settings between the manager and factory self.__synchronize_manager_factory_comm_settings() # Create data and logging dirs self.__create_data_and_logging_dirs() logger.debug("Starting TaskVineExecutor") # Create a process to run the TaskVine manager. submit_process_kwargs = {"ready_task_queue": self._ready_task_queue, "finished_task_queue": self._finished_task_queue, "should_stop": self._should_stop, "manager_config": self.manager_config} self._submit_process = multiprocessing.Process(target=_taskvine_submit_wait, name="TaskVine-Submit-Process", kwargs=submit_process_kwargs) # Create a process to run the TaskVine factory if enabled. if self.worker_launch_method == 'factory': factory_process_kwargs = {"should_stop": self._should_stop, "factory_config": self.factory_config} self._factory_process = multiprocessing.Process(target=_taskvine_factory, name="TaskVine-Factory-Process", kwargs=factory_process_kwargs) # Run thread to collect results and set tasks' futures. self._collector_thread = threading.Thread(target=self._collect_taskvine_results, name="TaskVine-Collector-Thread") # Interpreter can exit without waiting for this thread. self._collector_thread.daemon = True # Begin work self._submit_process.start() # Run worker scaler either with Parsl provider or TaskVine factory. # Skip if workers are launched manually. if self.worker_launch_method == 'factory': self._factory_process.start() elif self.worker_launch_method == 'provider': self.initialize_scaling() self._collector_thread.start() logger.debug("All components in TaskVineExecutor started") def _path_in_task(self, executor_task_id, *path_components): """ Returns a filename fixed and specific to a task. It is used for the following filename's: (not given): The subdirectory per task that contains function, result, etc. 'function': Pickled file that contains the function to be executed. 'argument': Pickled file that contains the arguments of the function call. 'result': Pickled file that (will) contain the result of the function. 'map': Pickled file with a dict between local parsl names, and remote taskvine names. """ task_dir = "{:04d}".format(executor_task_id) return os.path.join(self._function_data_dir, task_dir, *path_components) def submit(self, func, resource_specification, *args, **kwargs): """Processes the Parsl app by its arguments and submits the function information to the task queue, to be executed using the TaskVine system. The args and kwargs are processed for input and output files to the Parsl app, so that the files are appropriately specified for the TaskVine task. Parameters ---------- func : function Parsl app to be submitted to the TaskVine system resource_specification: dict Dictionary containing relevant info about task. Include information about resources of task, execution mode of task (out of {regular, serverless}). args : list Arguments to the Parsl app kwargs : dict Keyword arguments to the Parsl app """ logger.debug(f'Got resource specification: {resource_specification}') # Default execution mode of apps is regular exec_mode = resource_specification.get('exec_mode', self.function_exec_mode) # Detect resources and features of a submitted Parsl app cores = None memory = None disk = None gpus = None priority = None category = None running_time_min = None if resource_specification and isinstance(resource_specification, dict): for k in resource_specification: if k == 'cores': cores = resource_specification[k] elif k == 'memory': memory = resource_specification[k] elif k == 'disk': disk = resource_specification[k] elif k == 'gpus': gpus = resource_specification[k] elif k == 'priority': priority = resource_specification[k] elif k == 'category': category = resource_specification[k] elif k == 'running_time_min': running_time_min = resource_specification[k] # Assign executor task id to app executor_task_id = self._executor_task_counter self._executor_task_counter += 1 # Create a per task directory for the function, argument, map, and result files os.mkdir(self._path_in_task(executor_task_id)) input_files = [] output_files = [] # Determine whether to stage input files that will exist at the workers # Input and output files are always cached input_files += [self._register_file(f) for f in kwargs.get("inputs", []) if isinstance(f, File)] output_files += [self._register_file(f) for f in kwargs.get("outputs", []) if isinstance(f, File)] # Also consider any *arg that looks like a file as an input: input_files += [self._register_file(f) for f in args if isinstance(f, File)] for kwarg, maybe_file in kwargs.items(): # Add appropriate input and output files from "stdout" and "stderr" keyword arguments if kwarg == "stdout" or kwarg == "stderr": if maybe_file: output_files.append(self._std_output_to_vine(kwarg, maybe_file)) # For any other keyword that looks like a file, assume it is an input file elif isinstance(maybe_file, File): input_files.append(self._register_file(maybe_file)) # Create a Future object and have it be mapped from the task ID in the tasks dictionary fu = Future() fu.parsl_executor_task_id = executor_task_id with self._tasks_lock: self.tasks[str(executor_task_id)] = fu # Setup files to be used on a worker to execute the function function_file = None argument_file = None result_file = None map_file = None # Get path to files that will contain the pickled function, # arguments, result, and map of input and output files function_file = self._path_in_task(executor_task_id, "function") argument_file = self._path_in_task(executor_task_id, "argument") result_file = self._path_in_task(executor_task_id, "result") map_file = self._path_in_task(executor_task_id, "map") logger.debug("Creating executor task {} with function at: {}, argument at: {}, \ and result to be found at: {}".format(executor_task_id, function_file, argument_file, result_file)) # Serialize function object and arguments, separately self._serialize_object_to_file(function_file, func) args_dict = {'args': args, 'kwargs': kwargs} self._serialize_object_to_file(argument_file, args_dict) # Construct the map file of local filenames at worker self._construct_map_file(map_file, input_files, output_files) # Register a tarball containing all package dependencies for this app if instructed if self.manager_config.app_pack: env_pkg = self._prepare_package(func, self.extra_pkgs) else: env_pkg = None # Create message to put into the message queue logger.debug("Placing task {} on message queue".format(executor_task_id)) # Put apps into their categories based on function name if enabled if category is None: category = func.__name__ if self.manager_config.autocategory else 'parsl-default' task_info = ParslTaskToVine(executor_id=executor_task_id, exec_mode=exec_mode, category=category, input_files=input_files, output_files=output_files, map_file=map_file, function_file=function_file, argument_file=argument_file, result_file=result_file, cores=cores, memory=memory, disk=disk, gpus=gpus, priority=priority, running_time_min=running_time_min, env_pkg=env_pkg) # Send ready task to manager process if not self._submit_process.is_alive(): raise ExecutorError(self, "taskvine Submit Process is not alive") self._ready_task_queue.put_nowait(task_info) # Increment outstanding task counter with self._outstanding_tasks_lock: self._outstanding_tasks += 1 # Return the future for this function, will be set by collector thread when result # comes back from the TaskVine manager. return fu def _construct_worker_command(self): worker_command = self.factory_config.worker_executable if self.factory_config._project_password_file: worker_command += ' --password {}'.format(self.factory_config._project_password_file) if self.factory_config.worker_options: worker_command += ' {}'.format(self.factory_config.worker_options) if self.factory_config._project_name: worker_command += ' -M {}'.format(self.factory_config._project_name) else: worker_command += ' {} {}'.format(self.factory_config._project_address, self.factory_config._project_port) logger.debug("Using worker command: {}".format(worker_command)) return worker_command def _patch_providers(self): # Add the worker and password file to files that the provider needs to stage. # (Currently only for the CondorProvider) if isinstance(self.provider, CondorProvider): path_to_worker = shutil.which('vine_worker') self._worker_command = './' + self._worker_command self.provider.transfer_input_files.append(path_to_worker) if self.project_password_file: self.provider.transfer_input_files.append(self.project_password_file) def _serialize_object_to_file(self, path, obj): """Takes any object and serializes it to the file path.""" serialized_obj = serialize(obj, buffer_threshold=1024 * 1024) with open(path, 'wb') as f_out: written = 0 while written < len(serialized_obj): written += f_out.write(serialized_obj[written:]) def _construct_map_file(self, map_file, input_files, output_files): """ Map local filepath of parsl files to the filenames at the execution worker. If using a shared filesystem, the filepath is mapped to its absolute filename. Otherwise, to its original relative filename. In this later case, taskvine recreates any directory hierarchy needed.""" file_translation_map = {} for spec in itertools.chain(input_files, output_files): local_name = spec.parsl_name if self.manager_config.shared_fs: remote_name = os.path.abspath(local_name) else: remote_name = local_name file_translation_map[local_name] = remote_name self._serialize_object_to_file(map_file, file_translation_map) def _register_file(self, parsl_file): """Generates a tuple (parsl_file.filepath, stage, cache) to give to taskvine. cache is always True. stage is True if the file has a relative path. (i.e., not a URL or an absolute path)""" to_cache = True to_stage = False if parsl_file.scheme == 'file' or \ (parsl_file.local_path and os.path.exists(parsl_file.local_path)): to_stage = not os.path.isabs(parsl_file.filepath) return ParslFileToVine(parsl_file.filepath, to_stage, to_cache) def _std_output_to_vine(self, fdname, stdfspec): """Find the name of the file that will contain stdout or stderr and return a ParslFileToVine with it. These files are never cached""" fname, mode = putils.get_std_fname_mode(fdname, stdfspec) to_stage = not os.path.isabs(fname) return ParslFileToVine(fname, stage=to_stage, cache=False) def _prepare_package(self, fn, extra_pkgs): """ Look at source code of apps to figure out their package depedencies and output a tarball containing those to send along with tasks for execution.""" if not self._poncho_available: raise ExecutorError(self, 'poncho package is not available to individually pack apps.') fn_id = id(fn) fn_name = fn.__name__ if fn_id in self.cached_envs: logger.debug("Skipping analysis of %s, previously got %s", fn_name, self.cached_envs[fn_id]) return self.cached_envs[fn_id] source_code = inspect.getsource(fn).encode() pkg_dir = os.path.join(tempfile.gettempdir(), "python_package-{}".format(os.geteuid())) os.makedirs(pkg_dir, exist_ok=True) with tempfile.NamedTemporaryFile(suffix='.yaml') as spec: logger.info("Analyzing dependencies of %s", fn_name) analyze_cmdline = [self._package_analyze_script, exec_parsl_function.__file__, '-', spec.name] for p in extra_pkgs: analyze_cmdline += ["--extra-pkg", p] subprocess.run(analyze_cmdline, input=source_code, check=True) with open(spec.name, mode='rb') as f: spec_hash = hashlib.sha256(f.read()).hexdigest() logger.debug("Spec hash for %s is %s", fn_name, spec_hash) pkg = os.path.join(pkg_dir, "pack-{}.tar.gz".format(spec_hash)) if os.access(pkg, os.R_OK): self.cached_envs[fn_id] = pkg logger.debug("Cached package for %s found at %s", fn_name, pkg) return pkg (fd, tarball) = tempfile.mkstemp(dir=pkg_dir, prefix='.tmp', suffix='.tar.gz') os.close(fd) logger.info("Creating dependency package for %s", fn_name) logger.debug("Writing deps for %s to %s", fn_name, tarball) subprocess.run([self._package_create_script, spec.name, tarball], stdout=subprocess.DEVNULL, check=True) logger.debug("Done with conda-pack; moving %s to %s", tarball, pkg) os.rename(tarball, pkg) self.cached_envs[fn_id] = pkg return pkg def initialize_scaling(self): """ Compose the launch command and call scale out Scales the workers to the appropriate nodes with provider """ # Start scaling in/out logger.debug("Starting TaskVineExecutor with provider: %s", self.provider) self._worker_command = self._construct_worker_command() self._patch_providers() if hasattr(self.provider, 'init_blocks'): try: self.scale_out(blocks=self.provider.init_blocks) except Exception as e: logger.error("Initial block scaling out failed: {}".format(e)) raise e @property def outstanding(self) -> int: """Count the number of outstanding tasks.""" logger.debug(f"Counted {self._outstanding_tasks} outstanding tasks") return self._outstanding_tasks @property def workers_per_node(self) -> Union[int, float]: return 1 def scale_in(self, count): """Scale in method. Cancel a given number of blocks """ # Obtain list of blocks to kill to_kill = list(self.blocks.keys())[:count] kill_ids = [self.blocks[block] for block in to_kill] # Cancel the blocks provisioned if self.provider: self.provider.cancel(kill_ids) else: logger.error("No execution provider available to scale") def shutdown(self, *args, **kwargs): """Shutdown the executor. Sets flag to cancel the submit process and collector thread, which shuts down the TaskVine system submission. """ logger.debug("TaskVine shutdown started") self._should_stop.set() # Remove the workers that are still going kill_ids = [self.blocks[block] for block in self.blocks.keys()] if self.provider: logger.debug("Cancelling blocks") self.provider.cancel(kill_ids) # Join all processes before exiting logger.debug("Joining on submit process") self._submit_process.join() logger.debug("Joining on collector thread") self._collector_thread.join() if self.worker_launch_method == 'factory': logger.debug("Joining on factory process") self._factory_process.join() logger.debug("TaskVine shutdown completed") @wrap_with_logs def _collect_taskvine_results(self): """Sets the values of tasks' futures completed by taskvine. """ logger.debug("Starting Collector Thread") try: while not self._should_stop.is_set(): if not self._submit_process.is_alive(): raise ExecutorError(self, "taskvine Submit Process is not alive") # Get the result message from the _finished_task_queue try: task_report = self._finished_task_queue.get(timeout=1) except queue.Empty: continue # Obtain the future from the tasks dictionary with self._tasks_lock: future = self.tasks.pop(task_report.executor_id) logger.debug(f'Updating Future for Parsl Task: {task_report.executor_id}. \ Task {task_report.executor_id} has result_received set to {task_report.result_received}') if task_report.result_received: try: with open(task_report.result_file, 'rb') as f_in: result = deserialize(f_in.read()) except Exception as e: logger.error(f'Cannot load result from result file {task_report.result_file}. Exception: {e}') ex = TaskVineTaskFailure('Cannot load result from result file', None) ex.__cause__ = e future.set_exception(ex) else: if isinstance(result, Exception): ex = TaskVineTaskFailure('Task execution raises an exception', result) ex.__cause__ = result future.set_exception(ex) else: future.set_result(result) else: # If there are no results, then the task failed according to one of # taskvine modes, such as resource exhaustion. ex = TaskVineTaskFailure(task_report.reason, None) future.set_exception(ex) # decrement outstanding task counter with self._outstanding_tasks_lock: self._outstanding_tasks -= 1 finally: logger.debug(f"Marking all {self.outstanding} outstanding tasks as failed") logger.debug("Acquiring tasks_lock") with self._tasks_lock: logger.debug("Acquired tasks_lock") # set exception for tasks waiting for results that taskvine did not execute for fu in self.tasks.values(): if not fu.done(): fu.set_exception(TaskVineManagerFailure("taskvine executor failed to execute the task.")) logger.debug("Exiting Collector Thread")