From a185a120b99d0895e105d688036d062d982f7500 Mon Sep 17 00:00:00 2001 From: Sandro Wenzel Date: Sun, 20 Sep 2026 05:16:07 +0200 Subject: [PATCH] Add a modular workflow runner next to the existing one This is the main step of a refactoring of MC/bin/o2_dpg_workflow_runner.py, presented at CHEP 2026. The 2000-line script becomes a package, resource monitoring moves off the scheduling loop and costs about a tenth of what it did, and the scheduling policy becomes selectable, with two alternatives to the original one. Both runners are installed side by side and a dispatcher picks between them, so nothing changes for a caller that does not ask for the new one. - MC/bin/o2_dpg_workflow_runner.py becomes that dispatcher. It reads ALIEN_O2DPG_WORKFLOW_RUNNER and defaults to "legacy". The original runner moves unchanged to MC/bin/o2dpg_workflow_runner_legacy.py, and the o2dpg_workflow_runner.py symlink beside it is untouched, so all ten call sites in the repository keep working under either runner. - MC/workflow_runner/ holds the new package: workflow, graph, resources, monitoring, scheduler, executor, cleanup and cache modules. - Monitoring moves to a background thread with separate CPU and memory cadences. Polling psutil synchronously in the scheduling loop cost 10-20 % of a core on realistic workflows; this costs 1-2 %. - --scheduler-policy selects timeframe, which is the default and reproduces the original ordering, critical-path, or best-fit. - --systemd-run supersedes --cgroup for confining a workflow to a CPU and memory budget. - --cache-policy writes a fingerprint of command, environment, software tag and dependencies next to the _done marker, so a task whose command changed re-runs instead of being skipped. _done remains the skip marker. - Every flag the original parser accepts, the new one accepts. --cgroup, --webhook and --checkpoint-on-failure are accepted, ignored and warned about, so a JDL passing them through ALIEN_O2DPG_ADDITIONAL_WORKFLOW_RUNNER_ARGS still runs. - 68 unit tests come with it, run by a new CI job in .github/workflows/syntax-checks.yml. No Python test ran in CI before. - MC/workflow_runner/o2dpg_runner/README.md documents the layout, the behavioural differences and the pitfalls found while building it. https://indico.cern.ch/event/1471803/contributions/6967072/ Co-Authored-By: Claude Opus 5 --- .github/workflows/syntax-checks.yml | 15 + .gitignore | 1 + MC/bin/o2_dpg_workflow_runner.py | 2013 +---------------- MC/bin/o2dpg_workflow_runner_legacy.py | 2004 ++++++++++++++++ MC/workflow_runner/o2dpg_runner/README.md | 304 +++ MC/workflow_runner/o2dpg_runner/__init__.py | 11 + MC/workflow_runner/o2dpg_runner/alienv.py | 66 + MC/workflow_runner/o2dpg_runner/cache.py | 192 ++ MC/workflow_runner/o2dpg_runner/cleanup.py | 221 ++ MC/workflow_runner/o2dpg_runner/cli.py | 468 ++++ MC/workflow_runner/o2dpg_runner/config.py | 62 + MC/workflow_runner/o2dpg_runner/executor.py | 763 +++++++ MC/workflow_runner/o2dpg_runner/graph.py | 159 ++ MC/workflow_runner/o2dpg_runner/monitoring.py | 505 +++++ MC/workflow_runner/o2dpg_runner/resources.py | 362 +++ .../o2dpg_runner/scheduler/__init__.py | 22 + .../o2dpg_runner/scheduler/base.py | 62 + .../o2dpg_runner/scheduler/best_fit.py | 104 + .../o2dpg_runner/scheduler/critical_path.py | 60 + .../o2dpg_runner/scheduler/timeframe.py | 81 + .../o2dpg_runner/tests/__init__.py | 0 .../tests/fixtures/tiny_workflow.json | 104 + .../o2dpg_runner/tests/test_cache.py | 119 + .../o2dpg_runner/tests/test_executor_e2e.py | 216 ++ .../o2dpg_runner/tests/test_graph.py | 90 + .../o2dpg_runner/tests/test_resources.py | 137 ++ .../o2dpg_runner/tests/test_scheduler.py | 179 ++ .../o2dpg_runner/tests/test_simulator.py | 131 ++ .../o2dpg_runner/tests/test_workflow.py | 117 + MC/workflow_runner/o2dpg_runner/workflow.py | 411 ++++ .../o2dpg_schedule_simulator.py | 1268 +++++++++++ MC/workflow_runner/o2dpg_workflow_runner.py | 18 + 32 files changed, 8273 insertions(+), 1992 deletions(-) create mode 100755 MC/bin/o2dpg_workflow_runner_legacy.py create mode 100644 MC/workflow_runner/o2dpg_runner/README.md create mode 100644 MC/workflow_runner/o2dpg_runner/__init__.py create mode 100644 MC/workflow_runner/o2dpg_runner/alienv.py create mode 100644 MC/workflow_runner/o2dpg_runner/cache.py create mode 100644 MC/workflow_runner/o2dpg_runner/cleanup.py create mode 100644 MC/workflow_runner/o2dpg_runner/cli.py create mode 100644 MC/workflow_runner/o2dpg_runner/config.py create mode 100644 MC/workflow_runner/o2dpg_runner/executor.py create mode 100644 MC/workflow_runner/o2dpg_runner/graph.py create mode 100644 MC/workflow_runner/o2dpg_runner/monitoring.py create mode 100644 MC/workflow_runner/o2dpg_runner/resources.py create mode 100644 MC/workflow_runner/o2dpg_runner/scheduler/__init__.py create mode 100644 MC/workflow_runner/o2dpg_runner/scheduler/base.py create mode 100644 MC/workflow_runner/o2dpg_runner/scheduler/best_fit.py create mode 100644 MC/workflow_runner/o2dpg_runner/scheduler/critical_path.py create mode 100644 MC/workflow_runner/o2dpg_runner/scheduler/timeframe.py create mode 100644 MC/workflow_runner/o2dpg_runner/tests/__init__.py create mode 100644 MC/workflow_runner/o2dpg_runner/tests/fixtures/tiny_workflow.json create mode 100644 MC/workflow_runner/o2dpg_runner/tests/test_cache.py create mode 100644 MC/workflow_runner/o2dpg_runner/tests/test_executor_e2e.py create mode 100644 MC/workflow_runner/o2dpg_runner/tests/test_graph.py create mode 100644 MC/workflow_runner/o2dpg_runner/tests/test_resources.py create mode 100644 MC/workflow_runner/o2dpg_runner/tests/test_scheduler.py create mode 100644 MC/workflow_runner/o2dpg_runner/tests/test_simulator.py create mode 100644 MC/workflow_runner/o2dpg_runner/tests/test_workflow.py create mode 100644 MC/workflow_runner/o2dpg_runner/workflow.py create mode 100755 MC/workflow_runner/o2dpg_schedule_simulator.py create mode 100755 MC/workflow_runner/o2dpg_workflow_runner.py diff --git a/.github/workflows/syntax-checks.yml b/.github/workflows/syntax-checks.yml index 44bd0306f..942d7b1dd 100644 --- a/.github/workflows/syntax-checks.yml +++ b/.github/workflows/syntax-checks.yml @@ -105,6 +105,21 @@ jobs: done exit "$error" + runner-tests: + name: Workflow-runner unit tests + runs-on: ubuntu-latest + + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Install prerequisites + run: pip install pytest psutil + + - name: Run the o2dpg_runner test suite + working-directory: MC/workflow_runner + run: pytest o2dpg_runner/tests -q + pylint: name: Pylint runs-on: ubuntu-latest diff --git a/.gitignore b/.gitignore index 27281eba7..8271e2f98 100644 --- a/.gitignore +++ b/.gitignore @@ -2,3 +2,4 @@ .vscode *.pyc o2dpg_tests/** +__pycache__/ diff --git a/MC/bin/o2_dpg_workflow_runner.py b/MC/bin/o2_dpg_workflow_runner.py index eebde6501..fda79a19f 100755 --- a/MC/bin/o2_dpg_workflow_runner.py +++ b/MC/bin/o2_dpg_workflow_runner.py @@ -1,2004 +1,33 @@ #!/usr/bin/env python3 +"""Run the workflow runner selected by ALIEN_O2DPG_WORKFLOW_RUNNER. -# started February 2021, sandro.wenzel@cern.ch +'legacy' (the default) is the original single-file runner, 'new' is the +o2dpg_runner package under MC/workflow_runner. All arguments are passed on +unchanged. +""" -import re -import subprocess -import time -import json -import logging import os -import signal -import socket import sys -import traceback -import platform -import tarfile -from copy import deepcopy -try: - from graphviz import Digraph - havegraphviz=True -except ImportError: - havegraphviz=False - -formatter = logging.Formatter('%(asctime)s %(levelname)s %(message)s') - -sys.setrecursionlimit(100000) - -import argparse -import psutil -max_system_mem=psutil.virtual_memory().total - -sys.path.append(os.path.join(os.path.dirname(__file__), '.', 'o2dpg_workflow_utils')) -from o2dpg_workflow_utils import read_workflow - -# defining command line options -parser = argparse.ArgumentParser(description='Parallel execution of a (O2-DPG) DAG data/job pipeline under resource contraints.', - formatter_class=argparse.ArgumentDefaultsHelpFormatter) - -parser.add_argument('-f','--workflowfile', help='Input workflow file name', required=True) -parser.add_argument('-jmax','--maxjobs', type=int, help='Number of maximal parallel tasks.', default=100) -parser.add_argument('-k','--keep-going', action='store_true', help='Keep executing the pipeline as far possibe (not stopping on first failure)') -parser.add_argument('--dry-run', action='store_true', help='Show what you would do.') -parser.add_argument('--visualize-workflow', action='store_true', help='Saves a graph visualization of workflow.') -parser.add_argument('--target-labels', nargs='+', help='Runs the pipeline by target labels (example "TPC" or "DIGI").\ - This condition is used as logical AND together with --target-tasks.', default=[]) -parser.add_argument('-tt','--target-tasks', nargs='+', help='Runs the pipeline by target tasks (example "tpcdigi"). By default everything in the graph is run. Regular expressions supported.', default=["*"]) -parser.add_argument('--produce-script', help='Produces a shell script that runs the workflow in serialized manner and quits.') -parser.add_argument('--rerun-from', help='Reruns the workflow starting from given task (or pattern). All dependent jobs will be rerun.') -parser.add_argument('--list-tasks', help='Simply list all tasks by name and quit.', action='store_true') - -# Resources -parser.add_argument('--update-resources', dest="update_resources", help='Read resource estimates from a JSON and apply where possible.') -parser.add_argument("--dynamic-resources", dest="dynamic_resources", action="store_true", help="Update reources estimates of task based on finished related tasks") # derive resources dynamically -parser.add_argument('--optimistic-resources', dest="optimistic_resources", action="store_true", help="Try to run workflow even though resource limits might underestimate resource needs of some tasks") -parser.add_argument("--n-backfill", dest="n_backfill", type=int, default=1) -parser.add_argument('--mem-limit', help='Set memory limit as scheduling constraint (in MB)', default=0.9*max_system_mem/1024./1024, type=float) -parser.add_argument('--cpu-limit', help='Set CPU limit (core count)', default=8, type=float) -parser.add_argument('--cgroup', help='Execute pipeline under a given cgroup (e.g., 8coregrid) emulating resource constraints. This m\ -ust exist and the tasks file must be writable to with the current user.') - -# run control, webhooks -parser.add_argument('--stdout-on-failure', action='store_true', help='Print log files of failing tasks to stdout,') -parser.add_argument('--webhook', help=argparse.SUPPRESS) # log some infos to this webhook channel -parser.add_argument('--checkpoint-on-failure', help=argparse.SUPPRESS) # debug option making a debug-tarball and sending to specified address - # argument is alien-path -parser.add_argument('--retry-on-failure', help=argparse.SUPPRESS, default=0) # number of times a failing task is retried -parser.add_argument('--no-rootinit-speedup', help=argparse.SUPPRESS, action='store_true') # disable init of ROOT environment vars to speedup init/startup - -parser.add_argument('--remove-files-early', type=str, default="", help="Delete intermediate files early (using the file graph information in the given file)") - - -# Logging -parser.add_argument('--action-logfile', help='Logfilename for action logs. If none given, pipeline_action_#PID.log will be used') -parser.add_argument('--metric-logfile', help='Logfilename for metric logs. If none given, pipeline_metric_#PID.log will be used') -parser.add_argument('--production-mode', action='store_true', help='Production mode') -# will trigger special features good for non-interactive/production processing (automatic cleanup of files etc). -args = parser.parse_args() - -def setup_logger(name, log_file, level=logging.INFO): - """To setup as many loggers as you want""" - - handler = logging.FileHandler(log_file, mode='w') - handler.setFormatter(formatter) - - logger = logging.getLogger(name) - logger.setLevel(level) - logger.addHandler(handler) - - return logger - -# first file logger -actionlogger_file = ('pipeline_action_' + str(os.getpid()) + '.log', args.action_logfile)[args.action_logfile!=None] -actionlogger = setup_logger('pipeline_action_logger', actionlogger_file, level=logging.DEBUG) - -# second file logger -metriclogger = setup_logger('pipeline_metric_logger', ('pipeline_metric_' + str(os.getpid()) + '.log', args.action_logfile)[args.action_logfile!=None]) - -# Immediately log imposed memory and CPU limit as well as further useful meta info -_ , meta = read_workflow(args.workflowfile) -meta["cpu_limit"] = args.cpu_limit -meta["mem_limit"] = args.mem_limit -meta["workflow_file"] = os.path.abspath(args.workflowfile) -args.target_tasks = [f.strip('"').strip("'") for f in args.target_tasks] # strip quotes from the shell -meta["target_task"] = args.target_tasks -meta["rerun_from"] = args.rerun_from -meta["target_labels"] = args.target_labels -metriclogger.info(meta) - -# for debugging without terminal access -# TODO: integrate into standard logger -def send_webhook(hook, t): - if hook!=None: - command="curl -X POST -H 'Content-type: application/json' --data '{\"text\":\" " + str(t) + "\"}' " + str(hook) + " &> /dev/null" - os.system(command) - -# A fallback solution to getting all child procs -# in case psutil has problems (PermissionError). -# It returns the same list as psutil.children(recursive=True). -def getChildProcs(basepid): - cmd=''' - childprocs() { - local parent=$1 - if [ ! "$2" ]; then - child_pid_list="" - fi - if [ "$parent" ] ; then - child_pid_list="$child_pid_list $parent" - for childpid in $(pgrep -P ${parent}); do - childprocs $childpid "nottoplevel" - done; - fi - # return via a string list (only if toplevel) - if [ ! "$2" ]; then - echo "${child_pid_list}" - fi - } - ''' - cmd = cmd + '\n' + 'childprocs ' + str(basepid) - output = subprocess.check_output(cmd, shell=True) - plist = [] - for p in output.strip().split(): - try: - proc=psutil.Process(int(p)) - except psutil.NoSuchProcess: - continue - - plist.append(proc) - return plist - -# -# Code section to find all topological orderings -# of a DAG. This is used to know when we can schedule -# things in parallel. -# Taken from https://www.geeksforgeeks.org/all-topological-sorts-of-a-directed-acyclic-graph/ - -# class to represent a graph object -class Graph: - - # Constructor - def __init__(self, edges, N): - - # A List of Lists to represent an adjacency list - self.adjList = [[] for _ in range(N)] - - # stores in-degree of a vertex - # initialize in-degree of each vertex by 0 - self.indegree = [0] * N - - # add edges to the undirected graph - for (src, dest) in edges: - - # add an edge from source to destination - self.adjList[src].append(dest) - - # increment in-degree of destination vertex by 1 - self.indegree[dest] = self.indegree[dest] + 1 - -# Recursive function to find all topological orderings of a given DAG -def findAllTopologicalOrders(graph, path, discovered, N, allpaths, maxnumber=1): - if len(allpaths) >= maxnumber: - return - - # do for every vertex - for v in range(N): - - # proceed only if in-degree of current node is 0 and - # current node is not processed yet - if graph.indegree[v] == 0 and not discovered[v]: - - # for every adjacent vertex u of v, reduce in-degree of u by 1 - for u in graph.adjList[v]: - graph.indegree[u] = graph.indegree[u] - 1 - - # include current node in the path and mark it as discovered - path.append(v) - discovered[v] = True - - # recur - findAllTopologicalOrders(graph, path, discovered, N, allpaths) - - # backtrack: reset in-degree information for the current node - for u in graph.adjList[v]: - graph.indegree[u] = graph.indegree[u] + 1 - - # backtrack: remove current node from the path and - # mark it as undiscovered - path.pop() - discovered[v] = False - - # record valid ordering - if len(path) == N: - allpaths.append(path.copy()) - - -# get all topological orderings of a given DAG as a list -def printAllTopologicalOrders(graph, maxnumber=1): - # get number of nodes in the graph - N = len(graph.adjList) - - # create an auxiliary space to keep track of whether vertex is discovered - discovered = [False] * N - - # list to store the topological order - path = [] - allpaths = [] - # find all topological ordering and print them - findAllTopologicalOrders(graph, path, discovered, N, allpaths, maxnumber=maxnumber) - return allpaths - -# <--- end code section for topological sorts - -# find all tasks that depend on a given task (id); when a cache -# dict is given we can fill for the whole graph in one pass... -def find_all_dependent_tasks(possiblenexttask, tid, cache=None): - c=cache.get(tid) if cache else None - if c!=None: - return c - - daughterlist=[tid] - # possibly recurse - for n in possiblenexttask[tid]: - c = cache.get(n) if cache else None - if c == None: - c = find_all_dependent_tasks(possiblenexttask, n, cache) - daughterlist = daughterlist + c - if cache is not None: - cache[n]=c - - if cache is not None: - cache[tid]=daughterlist - return list(set(daughterlist)) - - -# wrapper taking some edges, constructing the graph, -# obtain all topological orderings and some other helper data structures -def analyseGraph(edges, nodes): - # Number of nodes in the graph - N = len(nodes) - - # candidate list trivial - nextjobtrivial = { n:[] for n in nodes } - # startnodes - nextjobtrivial[-1] = nodes - for e in edges: - nextjobtrivial[e[0]].append(e[1]) - if nextjobtrivial[-1].count(e[1]): - nextjobtrivial[-1].remove(e[1]) - - # find topological orderings of the graph - # create a graph from edges - graph = Graph(edges, N) - orderings = printAllTopologicalOrders(graph) - - return (orderings, nextjobtrivial) - - -def draw_workflow(workflowspec): - if not havegraphviz: - print('graphviz not installed, cannot draw workflow') - return - - dot = Digraph(comment='MC workflow') - nametoindex={} - index=0 - # nodes - for node in workflowspec['stages']: - name=node['name'] - nametoindex[name]=index - dot.node(str(index), name) - index=index+1 - - # edges - for node in workflowspec['stages']: - toindex = nametoindex[node['name']] - for req in node['needs']: - fromindex = nametoindex[req] - dot.edge(str(fromindex), str(toindex)) - - dot.render('workflow.gv') - -# builds the graph given a "taskuniverse" list -# builds accompagnying structures tasktoid and idtotask -def build_graph(taskuniverse, workflowspec): - tasktoid={ t[0]['name']:i for i, t in enumerate(taskuniverse, 0) } - # print (tasktoid) - - nodes = [] - edges = [] - for t in taskuniverse: - nodes.append(tasktoid[t[0]['name']]) - for n in t[0]['needs']: - edges.append((tasktoid[n], tasktoid[t[0]['name']])) - - return (edges, nodes) - - -# loads json into dict, e.g. for workflow specification -def load_json(workflowfile): - fp=open(workflowfile) - workflowspec=json.load(fp) - return workflowspec - - -# filters the original workflowspec according to wanted targets or labels -# returns a new workflowspec and the list of "final" workflowtargets -def filter_workflow(workflowspec, targets=[], targetlabels=[]): - if len(targets)==0: - return workflowspec, [] - if len(targetlabels)==0 and len(targets)==1 and targets[0]=="*": - return workflowspec, [] - - transformedworkflowspec = workflowspec - - def task_matches(t): - for filt in targets: - if filt=="*": - return True - if re.match(filt, t) != None: - return True - return False - - def task_matches_labels(t): - # when no labels are given at all it's ok - if len(targetlabels)==0: - return True - - for l in t['labels']: - if targetlabels.count(l)!=0: - return True - return False - - # The following sequence of operations works and is somewhat structured. - # However, it builds lookups used elsewhere as well, so some CPU might be saved by reusing - # some structures across functions or by doing less passes on the data. - - # helper lookup - tasknametoid = { t['name']:i for i, t in enumerate(workflowspec['stages'],0) } - - # check if a task can be run at all - # or not due to missing requirements - def canBeDone(t,cache={}): - ok = True - c = cache.get(t['name']) - if c != None: - return c - for r in t['needs']: - taskid = tasknametoid.get(r) - if taskid != None: - if not canBeDone(workflowspec['stages'][taskid], cache): - ok = False - break - else: - ok = False - break - cache[t['name']] = ok - if ok == False: - print (f"Disabling target {t['name']} due to unsatisfied requirements") - return ok - - okcache = {} - # build full target list - full_target_list = [ t for t in workflowspec['stages'] if task_matches(t['name']) and task_matches_labels(t) and canBeDone(t,okcache) ] - full_target_name_list = [ t['name'] for t in full_target_list ] - - # build full dependency list for a task t - def getallrequirements(t): - _l=[] - for r in t['needs']: - fulltask = workflowspec['stages'][tasknametoid[r]] - _l.append(fulltask) - _l=_l+getallrequirements(fulltask) - return _l - - full_requirements_list = [ getallrequirements(t) for t in full_target_list ] - - # make flat and fetch names only - full_requirements_name_list = list(set([ item['name'] for sublist in full_requirements_list for item in sublist ])) - - # inner "lambda" helper answering if a task "name" is needed by given targets - def needed_by_targets(name): - if full_target_name_list.count(name)!=0: - return True - if full_requirements_name_list.count(name)!=0: - return True - return False - - # we finaly copy everything matching the targets as well - # as all their requirements - transformedworkflowspec['stages']=[ l for l in workflowspec['stages'] if needed_by_targets(l['name']) ] - return transformedworkflowspec, full_target_name_list - - -# builds topological orderings (for each timeframe) -def build_dag_properties(workflowspec): - globaltaskuniverse = [ (l, i) for i, l in enumerate(workflowspec['stages'], 1) ] - timeframeset = set( l['timeframe'] for l in workflowspec['stages'] ) - - edges, nodes = build_graph(globaltaskuniverse, workflowspec) - tup = analyseGraph(edges, nodes.copy()) - # - global_next_tasks = tup[1] - - - dependency_cache = {} - # weight influences scheduling order can be anything user defined ... for the moment we just prefer to stay within a timeframe - # then take the number of tasks that depend on a task as further weight - # TODO: bring in resource estimates from runtime, CPU, MEM - # TODO: make this a policy of the runner to study different strategies - def getweight(tid): - return (globaltaskuniverse[tid][0]['timeframe'], len(find_all_dependent_tasks(global_next_tasks, tid, dependency_cache))) - - task_weights = [ getweight(tid) for tid in range(len(globaltaskuniverse)) ] - - for tid in range(len(globaltaskuniverse)): - actionlogger.info("Score for " + str(globaltaskuniverse[tid][0]['name']) + " is " + str(task_weights[tid])) - - # print (global_next_tasks) - return { 'nexttasks' : global_next_tasks, 'weights' : task_weights, 'topological_ordering' : tup[0] } - - -# update the resource estimates of a workflow based on resources given via JSON -def update_resource_estimates(workflow, resource_json): - # the resource_dict here is generated by tool o2dpg_sim_metrics.py json-stat - resource_dict = load_json(resource_json) - stages = workflow["stages"] - - for task in stages: - if task["timeframe"] >= 1: - name = "_".join(task["name"].split("_")[:-1]) - else: - name = task["name"] - - if name not in resource_dict: - continue - - new_resources = resource_dict[name] - - # memory - newmem = new_resources.get("pss", {}).get("max", None) - if newmem is not None: - oldmem = task["resources"]["mem"] - actionlogger.info("Updating mem estimate for " + task["name"] + " from " + str(oldmem) + " to " + str(newmem)) - task["resources"]["mem"] = newmem - - # cpu - newcpu = new_resources.get("cpu", {}).get("mean", None) - if newcpu is not None: - oldcpu = task["resources"]["cpu"] - rel_cpu = task["resources"]["relative_cpu"] - # TODO: No longer sure about this since we inject numbers from actually measured workloads - if rel_cpu is not None: - # respect the relative CPU settings - # By default, the CPU value in the workflow is already scaled if relative_cpu is given. - # The new estimate on the other hand is not yet scaled so it needs to be done here. - newcpu *= rel_cpu - actionlogger.info("Updating cpu estimate for " + task["name"] + " from " + str(oldcpu) + " to " + str(newcpu)) - task["resources"]["cpu"] = newcpu - -# a function to read a software environment determined by alienv into -# a python dictionary -def get_alienv_software_environment(packagestring): - """ - packagestring is something like O2::v202298081-1,O2Physics::xxx representing packages - published on CVMFS ... or ... a file containing directly the software environment to apply - """ - - # the trivial cases do nothing - if packagestring == None or packagestring == "" or packagestring == "None": - return {} - - def load_env_file(env_file): - """Transform an environment file generated with 'export > env.txt' into a python dictionary.""" - env_vars = {} - with open(env_file, "r") as f: - for line in f: - line = line.strip() - - # Ignore empty lines or comments - if not line or line.startswith("#"): - continue - - # Remove 'declare -x ' if present - if line.startswith("declare -x "): - line = line.replace("declare -x ", "", 1) - - # Handle case: "FOO" without "=" (assign empty string) - if "=" not in line: - key, value = line.strip(), "" - else: - key, value = line.split("=", 1) - value = value.strip('"') # Remove surrounding quotes if present - - env_vars[key.strip()] = value - return env_vars - - # see if this is a file - if os.path.exists(packagestring) and os.path.isfile(packagestring): - actionlogger.info("Taking software environment from file " + packagestring) - return load_env_file(packagestring) - - # alienv printenv packagestring --> dictionary - # for the moment this works with CVMFS only - cmd="/cvmfs/alice.cern.ch/bin/alienv printenv " + packagestring - proc = subprocess.Popen([cmd], stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True) - - envstring, err = proc.communicate() - # see if the printenv command was successful - if len(err.decode()) > 0: - print (err.decode()) - raise Exception - - # the software environment is now in the evnstring - # split it on semicolon - envstring=envstring.decode() - tokens=envstring.split(";") - # build envmap - envmap = {} - for t in tokens: - # check if assignment - if t.count("=") > 0: - assignment = t.rstrip().split("=") - envmap[assignment[0]] = assignment[1] - elif t.count("export") > 0: - # the case when we export or a simple variable - # need to consider the case when this has not been previously assigned - variable = t.split()[1] - if not variable in envmap: - envmap[variable]="" - - return envmap - -# -# functions for execution; encapsulated in a WorkflowExecutor class -# - -class Semaphore: - """ - Object that can be used as semaphore - """ - def __init__(self): - self.locked = False - def lock(self): - self.locked = True - def unlock(self): - self.locked = False - - -class ResourceBoundaries: - """ - Container holding global resource properties - """ - def __init__(self, cpu_limit, mem_limit, dynamic_resources=False, optimistic_resources=False): - self.cpu_limit = cpu_limit - self.mem_limit = mem_limit - self.dynamic_resources = dynamic_resources - # if this is set, tasks that would normally go beyond the resource limits will tried to be run in any case - self.optimistic_resources = optimistic_resources - - -class TaskResources: - """ - Container holding resources of a single task - """ - def __init__(self, tid, name, cpu, cpu_relative, mem, resource_boundaries): - # the task ID belonging to these resources - self.tid = tid - self.name = name - # original CPUs/MEM assigned (persistent) - self.cpu_assigned_original = cpu - self.mem_assigned_original = mem - # relative CPU, to be multiplied with sampled CPU; set by the user, e.g. to allow to backfill tasks - # only takes effect when sampling resources; persistent - self.cpu_relative = cpu_relative if cpu_relative else 1 - # CPUs/MEM assigned (transient) - self.cpu_assigned = cpu - self.mem_assigned = mem - # global resource settings - self.resource_boundaries = resource_boundaries - # sampled resources of this - self.cpu_sampled = None - self.mem_sampled = None - # Set these after a task has finished to compute new estimates for related tasks - self.walltime = None - self.cpu_taken = None - self.mem_taken = None - # collected during monitoring - self.time_collect = [] - self.cpu_collect = [] - self.mem_collect = [] - # linked to other resources of task that are of the same type as this one - self.related_tasks = None - # can assign a semaphore - self.semaphore = None - # the task's nice value - self.nice_value = None - # whether or not the task's resources are currently booked - self.booked = False - - @property - def is_done(self): - return self.time_collect and not self.booked - - def is_within_limits(self): - """ - Check if assigned resources respect limits - """ - cpu_within_limits = True - mem_within_limits = True - if self.cpu_assigned > self.resource_boundaries.cpu_limit: - cpu_within_limits = False - actionlogger.warning("CPU of task %s exceeds limits %d > %d", self.name, self.cpu_assigned, self.resource_boundaries.cpu_limit) - if self.cpu_assigned > self.resource_boundaries.mem_limit: - mem_within_limits = False - actionlogger.warning("MEM of task %s exceeds limits %d > %d", self.name, self.cpu_assigned, self.resource_boundaries.cpu_limit) - return cpu_within_limits and mem_within_limits - - def limit_resources(self, cpu_limit=None, mem_limit=None): - """ - Limit resources of this specific task - """ - if not cpu_limit: - cpu_limit = self.resource_boundaries.cpu_limit - if not mem_limit: - mem_limit = self.resource_boundaries.mem_limit - self.cpu_assigned = min(self.cpu_assigned, cpu_limit) - self.mem_assigned = min(self.mem_assigned, mem_limit) - - def add(self, time_passed, cpu, mem): - """ - Brief interface to add resources that were measured after time_passed - """ - self.time_collect.append(time_passed) - self.cpu_collect.append(cpu) - self.mem_collect.append(mem) - - def sample_resources(self): - """ - If this task is done, sample CPU and MEM for all related tasks that have not started yet - """ - if not self.is_done: - return - - if len(self.time_collect) < 3: - # Consider at least 3 points to sample from - self.cpu_sampled = self.cpu_assigned - self.mem_sampled = self.mem_assigned - actionlogger.debug("Task %s has not enough points (< 3) to sample resources, setting to previosuly assigned values.", self.name) - else: - # take the time deltas and leave out the very first CPU measurent which is not meaningful, - # at least when it domes from psutil.Proc.cpu_percent(interval=None) - time_deltas = [self.time_collect[i+1] - self.time_collect[i] for i in range(len(self.time_collect) - 1)] - cpu = sum([cpu * time_delta for cpu, time_delta in zip(self.cpu_collect[1:], time_deltas) if cpu >= 0]) - self.cpu_sampled = cpu / sum(time_deltas) - self.mem_sampled = max(self.mem_collect) - - mem_sampled = 0 - cpu_sampled = [] - for res in self.related_tasks: - if res.is_done: - mem_sampled = max(mem_sampled, res.mem_sampled) - cpu_sampled.append(res.cpu_sampled) - cpu_sampled = sum(cpu_sampled) / len(cpu_sampled) - - # This task ran already with the assigned resources, so let's set it to the limit - if cpu_sampled > self.resource_boundaries.cpu_limit: - actionlogger.warning("Sampled CPU (%.2f) exceeds assigned CPU limit (%.2f)", cpu_sampled, self.resource_boundaries.cpu_limit) - elif cpu_sampled < 0: - actionlogger.debug("Sampled CPU for %s is %.2f < 0, setting to previously assigned value %.2f", self.name, cpu_sampled, self.cpu_assigned) - cpu_sampled = self.cpu_assigned - - if mem_sampled > self.resource_boundaries.mem_limit: - actionlogger.warning("Sampled MEM (%.2f) exceeds assigned MEM limit (%.2f)", mem_sampled, self.resource_boundaries.mem_limit) - elif mem_sampled <= 0: - actionlogger.debug("Sampled memory for %s is %.2f <= 0, setting to previously assigned value %.2f", self.name, mem_sampled, self.mem_assigned) - mem_sampled = self.mem_assigned - - for res in self.related_tasks: - if res.is_done or res.booked: - continue - res.cpu_assigned = cpu_sampled * res.cpu_relative - res.mem_assigned = mem_sampled - # This task has been run before, stay optimistic and limit the resources in case the sampled ones exceed limits - res.limit_resources() - - -class ResourceManager: - """ - Central class to manage resources - - - CPU limits - - MEM limits - - Semaphores - - Entrypoint to set and to query for resources to be updated. - - Can be asked whether a certain task can be run under current resource usage. - Book and unbook resources. - """ - def __init__(self, cpu_limit, mem_limit, procs_parallel_max=100, dynamic_resources=False, optimistic_resources=False): - """ - Initialise members with defaults - """ - # hold TaskResources of all tasks - self.resources = [] - - # helper dictionaries holding common objects which will be distributed to single TaskResources objects - # to avoid further lookup and at the same time to share the same common objects - self.resources_related_tasks_dict = {} - self.semaphore_dict = {} - - # one common object that holds global resource settings such as CPU and MEM limits - self.resource_boundaries = ResourceBoundaries(cpu_limit, mem_limit, dynamic_resources, optimistic_resources) - - # register resources that are booked under default nice value - self.cpu_booked = 0 - self.mem_booked = 0 - # number of tasks currently booked - self.n_procs = 0 - - # register resources that are booked under high nice value - self.cpu_booked_backfill = 0 - self.mem_booked_backfill = 0 - # number of tasks currently booked under high nice value - self.n_procs_backfill = 0 - - # the maximum number of tasks that run at the same time - self.procs_parallel_max = procs_parallel_max - - # get the default nice value of this python script - self.nice_default = os.nice(0) - # add 19 to get nice value of low-priority tasks - self.nice_backfill = self.nice_default + 19 - - def add_task_resources(self, name, related_tasks_name, cpu, cpu_relative, mem, semaphore_string=None): - """ - Construct and Add a new TaskResources object - """ - resources = TaskResources(len(self.resources), name, cpu, cpu_relative, mem, self.resource_boundaries) - if not resources.is_within_limits() and not self.resource_boundaries.optimistic_resources: - # exit if we don't dare to try - print(f"Resources of task {name} are exceeding the boundaries.\nCPU: {cpu} (estimate) vs. {self.resource_boundaries.cpu_limit} (boundary)\nMEM: {mem} (estimated) vs. {self.resource_boundaries.mem_limit} (boundary).") - print("Pass --optimistic-resources to the runner to attempt the run anyway.") - exit(1) - # if we get here, either all is good or the user decided to be optimistic and we limit the resources, by default to the given CPU and mem limits. - resources.limit_resources() - - self.resources.append(resources) - # do the following to have the same Semaphore object for all corresponding TaskResources so that we do not need a lookup - if semaphore_string: - if semaphore_string not in self.semaphore_dict: - self.semaphore_dict[semaphore_string] = Semaphore() - resources.semaphore = self.semaphore_dict[semaphore_string] - - # do the following to give each TaskResources a list of the related tasks so we do not need an additional lookup - if related_tasks_name: - if related_tasks_name not in self.resources_related_tasks_dict: - # assigned list is [valid top be used, list of CPU, list of MEM, list of walltimes of each related task, list of processes that ran in parallel on average, list of taken CPUs, list of assigned CPUs, list of tasks finished in the meantime] - self.resources_related_tasks_dict[related_tasks_name] = [] - self.resources_related_tasks_dict[related_tasks_name].append(resources) - resources.related_tasks = self.resources_related_tasks_dict[related_tasks_name] - - def add_monitored_resources(self, tid, time_delta_since_start, cpu, mem): - self.resources[tid].add(time_delta_since_start, cpu, mem) - - def book(self, tid, nice_value): - """ - Book the resources of this task with given nice value - - The final nice value is determined by the final submission and could be different. - This can happen if the nice value should have been changed while that is not allowed by the system. - """ - res = self.resources[tid] - # take the nice value that was previously assigned when resources where checked last time - previous_nice_value = res.nice_value - - if previous_nice_value is None: - # this has not been checked ever if it was ok to be submitted - actionlogger.warning("Task ID %d has never been checked for resources. Treating as backfill", tid) - nice_value = self.nice_backfill - elif res.nice_value != nice_value: - actionlogger.warning("Task ID %d has was last time checked for a different nice value (%d) but is now submitted with (%d).", tid, res.nice_value, nice_value) - - res.nice_value = nice_value - res.booked = True - if res.semaphore is not None: - res.semaphore.lock() - if nice_value != self.nice_default: - self.n_procs_backfill += 1 - self.cpu_booked_backfill += res.cpu_assigned - self.mem_booked_backfill += res.mem_assigned - return - self.n_procs += 1 - self.cpu_booked += res.cpu_assigned - self.mem_booked += res.mem_assigned - - def unbook(self, tid): - """ - Unbook the reources of this task - """ - res = self.resources[tid] - res.booked = False - if self.resource_boundaries.dynamic_resources: - res.sample_resources() - if res.semaphore is not None: - res.semaphore.unlock() - if res.nice_value != self.nice_default: - self.cpu_booked_backfill -= res.cpu_assigned - self.mem_booked_backfill -= res.mem_assigned - self.n_procs_backfill -= 1 - if self.n_procs_backfill <= 0: - self.cpu_booked_backfill = 0 - self.mem_booked_backfill = 0 - return - self.n_procs -= 1 - self.cpu_booked -= res.cpu_assigned - self.mem_booked -= res.mem_assigned - if self.n_procs <= 0: - self.cpu_booked = 0 - self.mem_booked = 0 - - def ok_to_submit(self, tids): - """ - This generator yields the tid and nice value tuple from the list of task ids that should be checked - """ - tids_copy = tids.copy() - - def ok_to_submit_default(res): - """ - Return default nice value if conditions are met, None otherwise - """ - # analyse CPU - okcpu = (self.cpu_booked + res.cpu_assigned <= self.resource_boundaries.cpu_limit) - # analyse MEM - okmem = (self.mem_booked + res.mem_assigned <= self.resource_boundaries.mem_limit) - actionlogger.debug ('Condition check --normal-- for ' + str(res.tid) + ':' + res.name + ' CPU ' + str(okcpu) + ' MEM ' + str(okmem)) - return self.nice_default if (okcpu and okmem) else None - - def ok_to_submit_backfill(res, backfill_cpu_factor=1.5, backfill_mem_factor=1.5): - """ - Return backfill nice value if conditions are met, None otherwise - """ - if self.n_procs_backfill >= args.n_backfill: - return None - - if res.cpu_assigned > 0.9 * self.resource_boundaries.cpu_limit or res.mem_assigned / self.resource_boundaries.cpu_limit >= 1900: - return None - - # analyse CPU - okcpu = (self.cpu_booked_backfill + res.cpu_assigned <= self.resource_boundaries.cpu_limit) - okcpu = okcpu and (self.cpu_booked + self.cpu_booked_backfill + res.cpu_assigned <= backfill_cpu_factor * self.resource_boundaries.cpu_limit) - # analyse MEM - okmem = (self.mem_booked + self.mem_booked_backfill + res.mem_assigned <= backfill_mem_factor * self.resource_boundaries.mem_limit) - actionlogger.debug ('Condition check --backfill-- for ' + str(res.tid) + ':' + res.name + ' CPU ' + str(okcpu) + ' MEM ' + str(okmem)) - - return self.nice_backfill if (okcpu and okmem) else None - - if self.n_procs + self.n_procs_backfill >= self.procs_parallel_max: - # in this case, nothing can be done - return - - for ok_to_submit_impl, should_break in ((ok_to_submit_default, True), (ok_to_submit_backfill, False)): - tid_index = 0 - while tid_index < len(tids_copy): - - tid = tids_copy[tid_index] - res = self.resources[tid] - - actionlogger.info("Setup resources for task %s, cpu: %f, mem: %f", res.name, res.cpu_assigned, res.mem_assigned) - tid_index += 1 - - if (res.semaphore is not None and res.semaphore.locked) or res.booked: - continue - - nice_value = ok_to_submit_impl(res) - if nice_value is not None: - # if we get a non-None nice value, it means that this task is good to go - res.nice_value = nice_value - # yield the tid and its assigned nice value - yield tid, nice_value - - elif should_break: - # break here if resources of the next task do not fit - break - - -def filegraph_expand_timeframes(data: dict, timeframes: set, target_namelist) -> dict: - """ - A utility function for the fileaccess logic. Takes a template and duplicates - for the multi-timeframe structure. - """ - tf_entries = [ - entry for entry in data.get("file_report", []) - if re.match(r"^\./tf\d+/", entry["file"]) - ] - - result = {} - for i in timeframes: - if i == -1: - continue - # Deepcopy to avoid modifying original - new_entries = deepcopy(tf_entries) - for entry in new_entries: - # Fix filepath - entry["file"] = re.sub(r"^\./tf\d+/", f"./tf{i}/", entry["file"]) - # Fix written_by and read_by (preserve prefix, change numeric suffix) - entry["written_by"] = [ - re.sub(r"_\d+$", f"_{i}", w) for w in entry["written_by"] - ] - # for now we mark some files as keep if they are written - # by a target in the runner targetlist. TODO: Add other mechanisms - # to ask for file keeping (such as via regex or the like) - for e in entry["written_by"]: - if e in target_namelist: - entry["keep"] = True - entry["read_by"] = [ - re.sub(r"_\d+$", f"_{i}", r) for r in entry["read_by"] - ] - result[f"timeframe-{i}"] = new_entries - - return result - - - -class WorkflowExecutor: - # Constructor - def __init__(self, workflowfile, args, jmax=100): - self.args=args - self.is_productionmode = args.production_mode == True # os.getenv("ALIEN_PROC_ID") != None - self.workflowfile = workflowfile - self.workflowspec = load_json(workflowfile) - self.globalinit = self.extract_global_environment(self.workflowspec) # initialize global environment settings - for e in self.globalinit['env']: - if os.environ.get(e, None) == None: - value = self.globalinit['env'][e] - actionlogger.info("Applying global environment from init section " + str(e) + " : " + str(value)) - os.environ[e] = str(value) - - # only keep those tasks that are necessary to be executed based on user's filters - self.full_target_namelist = [] - self.workflowspec, self.full_target_namelist = filter_workflow(self.workflowspec, args.target_tasks, args.target_labels) - - if not self.workflowspec['stages']: - if args.target_tasks: - print ('Apparently some of the chosen target tasks are not in the workflow') - exit (0) - print ('Workflow is empty. Nothing to do') - exit (0) - - # construct the DAG, compute task weights - workflow = build_dag_properties(self.workflowspec) - if args.visualize_workflow: - draw_workflow(self.workflowspec) - self.possiblenexttask = workflow['nexttasks'] - self.taskweights = workflow['weights'] - self.topological_orderings = workflow['topological_ordering'] - self.taskuniverse = [ l['name'] for l in self.workflowspec['stages'] ] - # construct task ID <-> task name lookup - self.idtotask = [ 0 for _ in self.taskuniverse ] - self.tasktoid = {} - self.idtotf = [ l['timeframe'] for l in self.workflowspec['stages'] ] - for i, name in enumerate(self.taskuniverse): - self.tasktoid[name]=i - self.idtotask[i]=name - - if args.update_resources: - update_resource_estimates(self.workflowspec, args.update_resources) - - # construct the object that is in charge of resource management... - self.resource_manager = ResourceManager(args.cpu_limit, args.mem_limit, args.maxjobs, args.dynamic_resources, args.optimistic_resources) - for task in self.workflowspec['stages']: - # ...and add all initial resource estimates - global_task_name = self.get_global_task_name(task["name"]) - try: - cpu_relative = float(task["resources"]["relative_cpu"]) - except TypeError: - cpu_relative = 1 - self.resource_manager.add_task_resources(task["name"], global_task_name, float(task["resources"]["cpu"]), cpu_relative, float(task["resources"]["mem"]), task.get("semaphore")) - - self.procstatus = { tid:'ToDo' for tid in range(len(self.workflowspec['stages'])) } - self.taskneeds= { t:set(self.getallrequirements(t)) for t in self.taskuniverse } - self.stoponfailure = not args.keep_going - print ("Stop on failure ",self.stoponfailure) - - self.scheduling_iteration = 0 # count how often it was tried to schedule new tasks - self.process_list = [] # list of currently scheduled tasks with normal priority - self.backfill_process_list = [] # list of curently scheduled tasks with low backfill priority (not sure this is needed) - self.pid_to_psutilsproc = {} # cache of putilsproc for resource monitoring - self.pid_to_files = {} # we can auto-detect what files are produced by which task (at least to some extent) - self.pid_to_connections = {} # we can auto-detect what connections are opened by which task (at least to some extent) - signal.signal(signal.SIGINT, self.SIGHandler) - signal.siginterrupt(signal.SIGINT, False) - self.internalmonitorcounter = 0 # internal use - self.internalmonitorid = 0 # internal use - self.tids_marked_toretry = [] # sometimes we might want to retry a failed task (simply because it was "unlucky") and we put them here - self.retry_counter = [ 0 for tid in range(len(self.taskuniverse)) ] # we keep track of many times retried already - self.task_retries = [ self.workflowspec['stages'][tid].get('retry_count',0) for tid in range(len(self.taskuniverse)) ] # the per task specific "retry" number -> needs to be parsed from the JSON - - self.alternative_envs = {} # mapping of taskid to alternative software envs (to be applied on a per-task level) - # init alternative software environments - self.init_alternative_software_environments() - - # initialize container to keep track of file-task relationsships - self.file_removal_candidates = {} - self.do_early_file_removal = False - self.timeframeset = set([ task["timeframe"] for task in self.workflowspec['stages'] ]) - if args.remove_files_early != "": - with open(args.remove_files_early) as f: - filegraph_data = json.load(f) - self.do_early_file_removal = True - self.file_removal_candidates = filegraph_expand_timeframes(filegraph_data, self.timeframeset, self.full_target_namelist) - - def apply_global_env(self, environ_dict): - for e in self.globalinit['env']: - if environ_dict.get(e, None) == None: - value = self.globalinit['env'][e] - actionlogger.info("Applying global environment from init section " + str(e) + " : " + str(value)) - environ_dict[e] = str(value) - - def perform_early_file_removal(self, taskids): - """ - This function checks which files can be deleted upon completion of task - and optionally does so. - """ - - def remove_if_exists(filepath: str) -> None: - """ - Check if a file exists, and remove it if found. - """ - if os.path.exists(filepath): - fsize = os.path.getsize(filepath) - os.remove(filepath) - actionlogger.info(f"Removing {filepath} since no longer needed. Freeing {fsize/1024./1024.} MB.") - return True - - return False - - def remove_for_task_id(taskname, file_dict, timeframe_id, listofalltimeframes): - marked_for_removal = [] - - timeframestoscan = [ timeframe_id ] - if timeframe_id == -1: - timeframestoscan = [ i for i in listofalltimeframes if i != -1 ] - - # TODO: Note that this traversal of files is not certainly not optimal - # We should (and will) keep an mapping of tasks->potential files and just - # scan these. This is already provided by the FileIOGraph analysis tool. - for tid in timeframestoscan: - for i,file_entry in enumerate(file_dict[f"timeframe-{tid}"]): - filename = file_entry['file'] - read_by = file_entry['read_by'] - written_by = file_entry['written_by'] - if taskname in read_by: - file_entry['read_by'].remove(taskname) - if taskname in written_by: - file_entry['written_by'].remove(taskname) - - # TODO: in principle the written_by criterion might not be needed - if len(file_entry['read_by']) == 0 and len(file_entry['written_by']) == 0 and file_entry.get('keep', False) == False: - # the filename mentioned here is no longer needed and we can remove it - # make sure it is there and then delete it - if remove_if_exists(filename): - # also take out the file entry from the dict altogether - marked_for_removal.append(file_entry) - - #for k in marked_for_removal: - # file_dict[f"timeframe-{tid}"].remove(k) - - for tid in taskids: - taskname = self.idtotask[tid] - timeframe_id = self.idtotf[tid] - remove_for_task_id(taskname, self.file_removal_candidates, timeframe_id, self.timeframeset) - - - def SIGHandler(self, signum, frame): - """ - basically forcing shut down of all child processes - """ - actionlogger.info("Signal " + str(signum) + " caught") - try: - procs = psutil.Process().children(recursive=True) - except (psutil.NoSuchProcess): - pass - except (psutil.AccessDenied, PermissionError): - procs = getChildProcs(os.getpid()) - - for p in procs: - actionlogger.info("Terminating " + str(p)) - try: - p.terminate() - except (psutil.NoSuchProcess, psutil.AccessDenied): - pass - - _, alive = psutil.wait_procs(procs, timeout=3) - for p in alive: - try: - actionlogger.info("Killing " + str(p)) - p.kill() - except (psutil.NoSuchProcess, psutil.AccessDenied): - pass - - exit (1) - - def extract_global_environment(self, workflowspec): - """ - Checks if the workflow contains a dedicated init task - defining a global environment. Extract information and remove from workflowspec. - """ - init_index = 0 # this has to be the first task in the workflow - globalenv = {} - initcmd = None - if workflowspec['stages'][init_index]['name'] == '__global_init_task__': - env = workflowspec['stages'][init_index].get('env', None) - if env != None: - globalenv = { e : env[e] for e in env } - cmd = workflowspec['stages'][init_index].get('cmd', None) - if cmd != 'NO-COMMAND': - initcmd = cmd - - del workflowspec['stages'][init_index] - - return {"env" : globalenv, "cmd" : initcmd } - - def execute_globalinit_cmd(self, cmd): - actionlogger.info("Executing global setup cmd " + str(cmd)) - # perform the global init command (think of cleanup/setup things to be done in any case) - p = subprocess.Popen(['/bin/bash','-c', cmd], stdout=subprocess.PIPE, stderr=subprocess.PIPE) - stdout, stderr = p.communicate() - - # Check if the command was successful (return code 0) - if p.returncode == 0: - actionlogger.info(stdout.decode()) - else: - # this should be an error - actionlogger.error("Error executing global init function") - return False - return True - - def get_global_task_name(self, name): - """ - Get the global task name - - Tasks are related if only the suffix _ is different - """ - tokens = name.split("_") - try: - int(tokens[-1]) - return "_".join(tokens[:-1]) - except ValueError: - pass - return name - - def getallrequirements(self, task_name): - """ - get all requirement of a task by its name - """ - l=[] - for required_task_name in self.workflowspec['stages'][self.tasktoid[task_name]]['needs']: - l.append(required_task_name) - l=l+self.getallrequirements(required_task_name) - return l - - def get_logfile(self, tid): - """ - O2 taskwrapper logs task stdout and stderr to logfile .log - Get its exact path based on task ID - """ - # determines the logfile name for this task - name = self.workflowspec['stages'][tid]['name'] - workdir = self.workflowspec['stages'][tid]['cwd'] - return os.path.join(workdir, f"{name}.log") - - def get_done_filename(self, tid): - """ - O2 taskwrapper leaves .log_done after a task has successfully finished - Get its exact path based on task ID - """ - return f"{self.get_logfile(tid)}_done" - - def get_resources_filename(self, tid): - """ - O2 taskwrapper leaves .log_time after a task is done - Get its exact path based on task ID - """ - return f"{self.get_logfile(tid)}_time" - - # removes the done flag from tasks that need to be run again - def remove_done_flag(self, listoftaskids): - """ - Remove .log_done files to given task IDs - """ - for tid in listoftaskids: - done_filename = self.get_done_filename(tid) - name=self.workflowspec['stages'][tid]['name'] - if args.dry_run: - print ("Would mark task " + name + " as to be done again") - else: - print ("Marking task " + name + " as to be done again") - if os.path.exists(done_filename) and os.path.isfile(done_filename): - os.remove(done_filename) - - # submits a task as subprocess and records Popen instance - def submit(self, tid, nice): - """ - Submit a task - - 1. if needed, construct working directory if it does not yet exist - 2. update lookup structures flagging the task as being run - 3. set specific environment if requested for task - 4. construct psutil.Process from command line - 4.1 adjust the niceness of that process if requested - 5. return psutil.Process object - """ - actionlogger.debug("Submitting task " + str(self.idtotask[tid]) + " with nice value " + str(nice)) - c = self.workflowspec['stages'][tid]['cmd'] - workdir = self.workflowspec['stages'][tid]['cwd'] - if workdir: - if os.path.exists(workdir) and not os.path.isdir(workdir): - actionlogger.error('Cannot create working dir ... some other resource exists already') - return None - - if not os.path.isdir(workdir): - os.makedirs(workdir) - - self.procstatus[tid]='Running' - if args.dry_run: - drycommand="echo \' " + str(self.scheduling_iteration) + " : would do " + str(self.workflowspec['stages'][tid]['name']) + "\'" - return psutil.Popen(['/bin/bash','-c',drycommand], cwd=workdir) - - taskenv = os.environ.copy() - # apply specific (non-default) software version, if any - # (this was setup earlier) - alternative_env = self.alternative_envs.get(tid, None) - if alternative_env != None and len(alternative_env) > 0: - actionlogger.info('Applying alternative software environment to task ' + self.idtotask[tid]) - if alternative_env.get('TERM') != None: - # the environment is a complete environment - taskenv = {} - taskenv = alternative_env - else: - for entry in alternative_env: - # overwrite what is present in default - taskenv[entry] = alternative_env[entry] - - # add task specific environment - if self.workflowspec['stages'][tid].get('env')!=None: - taskenv.update(self.workflowspec['stages'][tid]['env']) - - # add global workflow environment - self.apply_global_env(taskenv) - - if os.environ.get('PIPELINE_RUNNER_DUMP_TASKENVS') != None: - envfilename = "taskenv_" + str(tid) + ".log" - with open(envfilename, "w") as file: - json.dump(taskenv, file, indent=2) - - p = psutil.Popen(['/bin/bash','-c',c], cwd=workdir, env=taskenv) - try: - p.nice(nice) - except (psutil.NoSuchProcess, psutil.AccessDenied): - actionlogger.error('Couldn\'t set nice value of ' + str(p.pid) + ' to ' + str(nice)) - - return p - - def ok_to_skip(self, tid): - """ - Decide if task can be skipped based on existence of .log_done - """ - done_filename = self.get_done_filename(tid) - if os.path.exists(done_filename) and os.path.isfile(done_filename): - return True - return False - - def try_job_from_candidates(self, taskcandidates, finished): - """ - Try to schedule next tasks - - Args: - taskcandidates: list - list of possible tasks that can be submitted - finished: list - empty list that will be filled with IDs of tasks that were finished in the meantime - """ - self.scheduling_iteration = self.scheduling_iteration + 1 - - # remove "done / skippable" tasks immediately - for tid in taskcandidates.copy(): # <--- the copy is important !! otherwise this loop is not doing what you think - if self.ok_to_skip(tid): - finished.append(tid) - taskcandidates.remove(tid) - actionlogger.info("Skipping task " + str(self.idtotask[tid])) - - # if tasks_skipped: - # return # ---> we return early in order to preserve some ordering (the next candidate tried should be daughters of skipped jobs) - # get task ID and proposed niceness from generator - for (tid, nice_value) in self.resource_manager.ok_to_submit(taskcandidates): - actionlogger.debug ("trying to submit " + str(tid) + ':' + str(self.idtotask[tid])) - if p := self.submit(tid, nice_value): - # explicitly set the nice value here from the process again because it might happen that submit could not change the niceness - # so we let the ResourceManager know what the final niceness is - self.resource_manager.book(tid, p.nice()) - self.process_list.append((tid,p)) - taskcandidates.remove(tid) - # minimal delay - time.sleep(0.1) - - def stop_pipeline_and_exit(self, process_list): - # kill all remaining jobs - for p in process_list: - p[1].kill() - - exit(1) - - - def monitor(self, process_list): - """ - Go through all running tasks and get their current resources - - Resources are summed up for tasks and all their children - - Pass CPU, PSS, USS, niceness, current time to metriclogger - - Warn if overall PSS exceeds assigned memory limit - """ - self.internalmonitorcounter+=1 - if self.internalmonitorcounter % 5 != 0: - return - - self.internalmonitorid+=1 - - globalCPU=0. - globalPSS=0. - resources_per_task = {} - - # On a global level, we are interested in total disc space used (not differential in tasks) - # We can call system "du" as the fastest impl - def disk_usage_du(path: str) -> int: - """Use system du to get total size in bytes.""" - out = subprocess.check_output(['du', '-sb', path], text=True) - return int(out.split()[0]) - - disc_usage = -1 - if os.getenv("MONITOR_DISC_USAGE"): - disc_usage = disk_usage_du(os.getcwd()) / 1024. / 1024 # in MB - - for tid, proc in process_list: - - # proc is Popen object - pid=proc.pid - if self.pid_to_files.get(pid)==None: - self.pid_to_files[pid]=set() - self.pid_to_connections[pid]=set() - try: - psutilProcs = [ proc ] - # use psutil for CPU measurement - psutilProcs = psutilProcs + proc.children(recursive=True) - except (psutil.NoSuchProcess): - continue - - except (psutil.AccessDenied, PermissionError): - psutilProcs = psutilProcs + getChildProcs(pid) - - # accumulate total metrics (CPU, memory) - totalCPU = 0. - totalPSS = 0. - totalSWAP = 0. - totalUSS = 0. - for p in psutilProcs: - """ - try: - for f in p.open_files(): - self.pid_to_files[pid].add(str(f.path)+'_'+str(f.mode)) - for f in p.connections(kind="all"): - remote=f.raddr - if remote==None: - remote='none' - self.pid_to_connections[pid].add(str(f.type)+"_"+str(f.laddr)+"_"+str(remote)) - except Exception: - pass - """ - thispss=0 - thisuss=0 - # MEMORY part - try: - fullmem=p.memory_full_info() - thispss=getattr(fullmem,'pss',0) #<-- pss not available on MacOS - totalPSS=totalPSS + thispss - totalSWAP=totalSWAP + fullmem.swap - thisuss=fullmem.uss - totalUSS=totalUSS + thisuss - except (psutil.NoSuchProcess, psutil.AccessDenied): - pass - - # CPU part - # fetch existing proc or insert - cachedproc = self.pid_to_psutilsproc.get(p.pid) - if cachedproc!=None: - try: - thiscpu = cachedproc.cpu_percent(interval=None) - except (psutil.NoSuchProcess, psutil.AccessDenied): - thiscpu = 0. - totalCPU = totalCPU + thiscpu - # thisresource = {'iter':self.internalmonitorid, 'pid': p.pid, 'cpu':thiscpu, 'uss':thisuss/1024./1024., 'pss':thispss/1024./1024.} - # metriclogger.info(thisresource) - else: - self.pid_to_psutilsproc[p.pid] = p - try: - self.pid_to_psutilsproc[p.pid].cpu_percent() - except (psutil.NoSuchProcess, psutil.AccessDenied): - pass - - time_delta = int((time.perf_counter() - self.start_time) * 1000) - totalUSS = totalUSS / 1024 / 1024 - totalPSS = totalPSS / 1024 / 1024 - nice_value = proc.nice() - resources_per_task[tid]={'iter':self.internalmonitorid, - 'name':self.idtotask[tid], - 'cpu':totalCPU, - 'uss':totalUSS, - 'pss':totalPSS, - 'nice':nice_value, - 'swap':totalSWAP, - 'label':self.workflowspec['stages'][tid]['labels'], - 'disc': disc_usage} - self.resource_manager.add_monitored_resources(tid, time_delta, totalCPU / 100, totalPSS) - if nice_value == self.resource_manager.nice_default: - globalCPU += totalCPU - globalPSS += totalPSS - - metriclogger.info(resources_per_task[tid]) - send_webhook(self.args.webhook, resources_per_task) - - if globalPSS > self.resource_manager.resource_boundaries.mem_limit: - metriclogger.info('*** MEMORY LIMIT PASSED !! ***') - # --> We could use this for corrective actions such as killing jobs currently back-filling - # (or better hibernating) - - def waitforany(self, process_list, finished, failingtasks): - """ - Loop through all submitted tasks and check if they are finished - - 1. If process is still running, do nothing - 2. If process is finished, get its return value, update finished and failingtasks lists - 2.1 unbook resources - 2.2 add taken resources and pass the to ResourceManager - """ - failuredetected = False - failingpids = [] - if len(process_list)==0: - return False - - for p in list(process_list): - pid = p[1].pid - tid = p[0] # the task id of this process - returncode = 0 - if not self.args.dry_run: - returncode = p[1].poll() - if returncode!=None: - actionlogger.info ('Task ' + str(pid) + ' ' + str(tid)+':'+str(self.idtotask[tid]) + ' finished with status ' + str(returncode)) - # account for cleared resources - self.resource_manager.unbook(tid) - self.procstatus[tid]='Done' - finished.append(tid) - #self.validate_resources_running(tid) - process_list.remove(p) - if returncode != 0: - print (str(self.idtotask[tid]) + ' failed ... checking retry') - # we inspect if this is something "unlucky" which could be resolved by a simple resubmit - if self.is_worth_retrying(tid) and ((self.retry_counter[tid] < int(args.retry_on_failure)) or (self.retry_counter[tid] < int(self.task_retries[tid]))): - print (str(self.idtotask[tid]) + ' to be retried') - actionlogger.info ('Task ' + str(self.idtotask[tid]) + ' failed but marked to be retried ') - self.tids_marked_toretry.append(tid) - self.retry_counter[tid] += 1 - - else: - failuredetected = True - failingpids.append(pid) - failingtasks.append(tid) - - if failuredetected and self.stoponfailure: - actionlogger.info('Stoping pipeline due to failure in stages with PID ' + str(failingpids)) - # self.analyse_files_and_connections() - if self.args.stdout_on_failure: - self.cat_logfiles_tostdout(failingtasks) - self.send_checkpoint(failingtasks, self.args.checkpoint_on_failure) - self.stop_pipeline_and_exit(process_list) - - # empty finished means we have to wait more - return len(finished)==0 - - def is_worth_retrying(self, tid): - # This checks for some signatures in logfiles that indicate that a retry of this task - # might have a chance. - # Ideally, this should be made user configurable. Either the user could inject a lambda - # or a regular expression to use. For now we just put a hard coded list - logfile = self.get_logfile(tid) - - return True #! --> for now we just retry tasks a few times - - # 1) ZMQ_EVENT + interrupted system calls (DPL bug during shutdown) - # Not sure if grep is faster than native Python text search ... - # status = os.system('grep "failed setting ZMQ_EVENTS" ' + logfile + ' &> /dev/null') - # if os.WEXITSTATUS(status) == 0: - # return True - - # return False - - - def cat_logfiles_tostdout(self, taskids): - # In case of errors we can cat the logfiles for this taskname - # to stdout. Assuming convention that "taskname" translates to "taskname.log" logfile. - for tid in taskids: - logfile = self.get_logfile(tid) - if os.path.exists(logfile): - print (' ----> START OF LOGFILE ', logfile, ' -----') - os.system('cat ' + logfile) - print (' <---- END OF LOGFILE ', logfile, ' -----') - - def send_checkpoint(self, taskids, location): - # Makes a tarball containing all files in the base dir - # (timeframe independent) and the dir with corrupted timeframes - # and copies it to a specific ALIEN location. Not a core function - # just some tool get hold on error conditions appearing on the GRID. - - def get_tar_command(dir='./', flags='cf', findtype='f', filename='checkpoint.tar'): - return 'find ' + str(dir) + ' -maxdepth 1 -type ' + str(findtype) + ' -print0 | xargs -0 tar ' + str(flags) + ' ' + str(filename) - - if location != None: - print ('Making a failure checkpoint') - # let's determine a filename from ALIEN_PROC_ID - hostname - and PID - - aliprocid=os.environ.get('ALIEN_PROC_ID') - if aliprocid == None: - aliprocid = 0 - - fn='pipeline_checkpoint_ALIENPROC' + str(aliprocid) + '_PID' + str(os.getpid()) + '_HOST' + socket.gethostname() + '.tar' - actionlogger.info("Checkpointing to file " + fn) - tarcommand = get_tar_command(filename=fn) - actionlogger.info("Taring " + tarcommand) - - # create a README file with instruction on how to use checkpoint - readmefile=open('README_CHECKPOINT_PID' + str(os.getpid()) + '.txt','w') - - for tid in taskids: - taskspec = self.workflowspec['stages'][tid] - name = taskspec['name'] - readmefile.write('Checkpoint created because of failure in task ' + name + '\n') - readmefile.write('In order to reproduce with this checkpoint, do the following steps:\n') - readmefile.write('a) setup the appropriate O2sim environment using alienv\n') - readmefile.write('b) run: $O2DPG_ROOT/MC/bin/o2_dpg_workflow_runner.py -f workflow.json -tt ' + name + '$ --retry-on-failure 0\n') - readmefile.close() - - # first of all the base directory - os.system(tarcommand) - - # then we add stuff for the specific timeframes ids if any - for tid in taskids: - taskspec = self.workflowspec['stages'][tid] - directory = taskspec['cwd'] - if directory != "./": - tarcommand = get_tar_command(dir=directory, flags='rf', filename=fn) - actionlogger.info("Tar command is " + tarcommand) - os.system(tarcommand) - # same for soft links - tarcommand = get_tar_command(dir=directory, flags='rf', findtype='l', filename=fn) - actionlogger.info("Tar command is " + tarcommand) - os.system(tarcommand) - - # prepend file:/// to denote local file - fn = "file://" + fn - actionlogger.info("Local checkpoint file is " + fn) - - # location needs to be an alien path of the form alien:///foo/bar/ - copycommand='alien.py cp ' + fn + ' ' + str(location) + '@disk:1' - actionlogger.info("Copying to alien " + copycommand) - os.system(copycommand) - - def init_alternative_software_environments(self): - """ - Initialises alternative software environments for specific tasks, if there - is an annotation in the workflow specificiation. - """ - - environment_cache = {} - # go through all the tasks once and setup environment - for taskid in range(len(self.workflowspec['stages'])): - packagestr = self.workflowspec['stages'][taskid].get("alternative_alienv_package") - if packagestr == None: - continue - - if environment_cache.get(packagestr) == None: - environment_cache[packagestr] = get_alienv_software_environment(packagestr) - - self.alternative_envs[taskid] = environment_cache[packagestr] - - - def analyse_files_and_connections(self): - for p,s in self.pid_to_files.items(): - for f in s: - print("F" + str(f) + " : " + str(p)) - for p,s in self.pid_to_connections.items(): - for c in s: - print("C" + str(c) + " : " + str(p)) - #print(str(p) + " CONS " + str(c)) - try: - # check for intersections - for p1, s1 in self.pid_to_files.items(): - for p2, s2 in self.pid_to_files.items(): - if p1!=p2: - if type(s1) is set and type(s2) is set: - if len(s1)>0 and len(s2)>0: - try: - inters = s1.intersection(s2) - except Exception: - print ('Exception during intersect inner') - pass - if (len(inters)>0): - print ('FILE Intersection ' + str(p1) + ' ' + str(p2) + ' ' + str(inters)) - # check for intersections - for p1, s1 in self.pid_to_connections.items(): - for p2, s2 in self.pid_to_connections.items(): - if p1!=p2: - if type(s1) is set and type(s2) is set: - if len(s1)>0 and len(s2)>0: - try: - inters = s1.intersection(s2) - except Exception: - print ('Exception during intersect inner') - pass - if (len(inters)>0): - print ('CON Intersection ' + str(p1) + ' ' + str(p2) + ' ' + str(inters)) - - # check for intersections - #for p1, s1 in slf.pid_to_files.items(): - # for p2, s2 in self.pid_to_files.items(): - # if p1!=p2 and len(s1.intersection(s2))!=0: - # print ('Intersection found files ' + str(p1) + ' ' + str(p2) + ' ' + s1.intersection(s2)) - except Exception as e: - exc_type, exc_obj, exc_tb = sys.exc_info() - fname = os.path.split(exc_tb.tb_frame.f_code.co_filename)[1] - print(exc_type, fname, exc_tb.tb_lineno) - print('Exception during intersect outer') - pass - - def is_good_candidate(self, candid, finishedtasks): - if self.procstatus[candid] != 'ToDo': - return False - needs = set([self.tasktoid[t] for t in self.taskneeds[self.idtotask[candid]]]) - if set(finishedtasks).intersection(needs) == needs: - return True - return False - - def emit_code_for_task(self, tid, lines): - actionlogger.debug("Submitting task " + str(self.idtotask[tid])) - taskspec = self.workflowspec['stages'][tid] - c = taskspec['cmd'] - workdir = taskspec['cwd'] - env = taskspec.get('env') - # in general: - # try to make folder - lines.append('[ ! -d ' + workdir + ' ] && mkdir ' + workdir + '\n') - # cd folder - lines.append('cd ' + workdir + '\n') - # set local environment - if env!=None: - for e in env.items(): - lines.append('export ' + e[0] + '=' + str(e[1]) + '\n') - # do command - lines.append(c + '\n') - # unset local environment - if env!=None: - for e in env.items(): - lines.append('unset ' + e[0] + '\n') - - # cd back - lines.append('cd $OLDPWD\n') - - - # produce a bash script that runs workflow standalone - def produce_script(self, filename): - # pick one of the correct task orderings - taskorder = self.topological_orderings[0] - outF = open(filename, "w") - - lines=[] - # header - lines.append('#!/usr/bin/env bash\n') - lines.append('#THIS FILE IS AUTOGENERATED\n') - lines.append('export JOBUTILS_SKIPDONE=ON\n') - - # we record the global environment setting - # in particular to capture global workflow initialization - lines.append('#-- GLOBAL INIT SECTION FROM WORKFLOW --\n') - for e in self.globalinit['env']: - lines.append('export ' + str(e) + '=' + str(self.globalinit['env'][e]) + '\n') - lines.append('#-- TASKS FROM WORKFLOW --\n') - for tid in taskorder: - print ('Doing task ' + self.idtotask[tid]) - self.emit_code_for_task(tid, lines) - - outF.writelines(lines) - outF.close() - - def production_endoftask_hook(self, tid): - # Executes a hook at end of a successful task, meant to be used in GRID productions. - # For the moment, archiving away log files, done + time files from jobutils. - # TODO: In future this may be much more generic tasks such as dynamic cleanup of intermediate - # files (when they are no longer needed). - # TODO: Care must be taken with the continue feature as `_done` files are stored elsewhere now - actionlogger.info("Cleaning up log files for task " + str(tid)) - logf = self.get_logfile(tid) - donef = self.get_done_filename(tid) - timef = logf + "_time" - - # add to tar file archive - tf = tarfile.open(name="pipeline_log_archive.log.tar", mode='a') - if tf != None: - tf.add(logf) - tf.add(donef) - tf.add(timef) - tf.close() - - # remove original file - os.remove(logf) - os.remove(donef) - os.remove(timef) - - # print error message when no progress can be made - def noprogress_errormsg(self): - # TODO: rather than writing this out here; refer to the documentation discussion this? - msg = """Scheduler runtime error: The scheduler is not able to make progress although we have a non-zero candidate set. - -Explanation: This is typically the case because the **ESTIMATED** resource requirements for some tasks -in the workflow exceed the available number of CPU cores or the memory (as explicitely or implicitely determined from the ---cpu-limit and --mem-limit options). Often, this might be the case on laptops with <=16GB of RAM if one of the tasks -is demanding ~16GB. In this case, one could try to tell the scheduler to use a slightly higher memory limit -with an explicit --mem-limit option (for instance `--mem-limit 20000` to set to 20GB). This might work whenever the -**ACTUAL** resource usage of the tasks is smaller than anticipated (because only small test cases are run). - -In addition it might be worthwile running the workflow without this resource aware, dynamic scheduler. -This is possible by converting the json workflow into a linearized shell script and by directly executing the shell script. -Use the `--produce-script myscript.sh` option for this. -""" - print (msg, file=sys.stderr) - - def execute(self): - self.start_time = time.perf_counter() - psutil.cpu_percent(interval=None) - os.environ['JOBUTILS_SKIPDONE'] = "ON" - errorencountered = False - - def speedup_ROOT_Init(): - """initialize some env variables that speed up ROOT init - and prevent ROOT from spawning many short-lived child - processes""" - - # only do it on Linux - if platform.system() != 'Linux': - return - - if os.environ.get('ROOT_LDSYSPATH')!=None and os.environ.get('ROOT_CPPSYSINCL')!=None: - # do nothing if already defined - return - - # a) the PATH for system libraries - # search taken from ROOT TUnixSystem - cmd='LD_DEBUG=libs LD_PRELOAD=DOESNOTEXIST ls /tmp/DOESNOTEXIST 2>&1 | grep -m 1 "system search path" | sed \'s/.*=//g\' | awk \'//{print $1}\'' - proc = subprocess.Popen([cmd], stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True) - libpath, err = proc.communicate() - if not (args.no_rootinit_speedup == True): - print ("setting up ROOT system") - os.environ['ROOT_LDSYSPATH'] = libpath.decode() - os.environ['CLING_LDSYSPATH'] = libpath.decode() - - # b) the PATH for compiler includes needed by Cling - cmd = "LC_ALL=C c++ -xc++ -E -v /dev/null 2>&1 | sed -n '/^#include/,${/^ \\/.*++/{p}}'" - proc = subprocess.Popen([cmd], stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True) - incpath, err = proc.communicate() - incpaths = [ line.lstrip() for line in incpath.decode().splitlines() ] - joined = ':'.join(incpaths) - if not (args.no_rootinit_speedup == True): - actionlogger.info("Determined ROOT_CPPSYSINCL=" + joined) - os.environ['ROOT_CPPSYSINCL'] = joined - os.environ['CLING_CPPSYSINCL'] = joined - - speedup_ROOT_Init() - - # we make our own "tmp" folder - # where we can put stuff such as tmp socket files etc (for instance DPL FAIR-MQ sockets) - # (In case of running within docker/singularity, this may not be so important) - if not os.path.isdir("./.tmp"): - os.mkdir("./.tmp") - if os.environ.get('FAIRMQ_IPC_PREFIX')==None: - socketpath = os.getcwd() + "/.tmp" - actionlogger.info("Setting FAIRMQ socket path to " + socketpath) - os.environ['FAIRMQ_IPC_PREFIX'] = socketpath - - # some maintenance / init work - if args.list_tasks: - print ('List of tasks in this workflow:') - for i,t in enumerate(self.workflowspec['stages'],0): - print (t['name'] + ' (' + str(t['labels']) + ')' + ' ToDo: ' + str(not self.ok_to_skip(i))) - exit (0) - - if args.produce_script != None: - self.produce_script(args.produce_script) - exit (0) - - # execute the user-given global init cmd for this workflow - globalinitcmd = self.globalinit.get("cmd", None) - if globalinitcmd != None: - if not self.execute_globalinit_cmd(globalinitcmd): - exit (1) - - if args.rerun_from: - reruntaskfound=False - for task in self.workflowspec['stages']: - taskname=task['name'] - if re.match(args.rerun_from, taskname): - reruntaskfound=True - taskid=self.tasktoid[taskname] - self.remove_done_flag(find_all_dependent_tasks(self.possiblenexttask, taskid)) - if not reruntaskfound: - print('No task matching ' + args.rerun_from + ' found; cowardly refusing to do anything ') - exit (1) - - # ***************** - # main control loop - # ***************** - candidates = [ tid for tid in self.possiblenexttask[-1] ] - - self.process_list=[] # list of tuples of nodes ids and Popen subprocess instances - - finishedtasks=[] # global list of finished tasks - - try: - - while True: - # sort candidate list according to task weights - candidates = [ (tid, self.taskweights[tid]) for tid in candidates ] - candidates.sort(key=lambda tup: (tup[1][0],-tup[1][1])) # prefer small and same timeframes first then prefer important tasks within frameframe - # remove weights - candidates = [ tid for tid,_ in candidates ] - - finished = [] # --> to account for finished because already done or skipped - actionlogger.debug('Sorted current candidates: ' + str([(c,self.idtotask[c]) for c in candidates])) - self.try_job_from_candidates(candidates, finished) - if len(candidates) > 0 and len(self.process_list) == 0: - self.noprogress_errormsg() - send_webhook(self.args.webhook,"Unable to make further progress: Quitting") - errorencountered = True - break - - finished_from_started = [] # to account for finished when actually started - failing = [] - while self.waitforany(self.process_list, finished_from_started, failing): - if not args.dry_run: - self.monitor(self.process_list) # ---> make async to normal operation? - time.sleep(1) # <--- make this incremental (small wait at beginning) - else: - time.sleep(0.001) - - finished = finished + finished_from_started - actionlogger.debug("finished now :" + str(finished_from_started)) - finishedtasks = finishedtasks + finished - - # perform file cleanup - if self.do_early_file_removal: - self.perform_early_file_removal(finished_from_started) - - if self.is_productionmode: - # we can do some generic cleanup of finished tasks in non-interactive/GRID mode - # TODO: this can run asynchronously - for _t in finished_from_started: - self.production_endoftask_hook(_t) - - # if a task was marked "failed" and we come here (because - # we use --keep-going) ... we need to take out the pid from finished - if len(failing) > 0: - # remove these from those marked finished in order - # not to continue with their children - errorencountered = True - for t in failing: - finished = [ x for x in finished if x != t ] - finishedtasks = [ x for x in finishedtasks if x != t ] - - # if a task was marked as "retry" we simply put it back into the candidate list - if len(self.tids_marked_toretry) > 0: - # we need to remove these first of all from those marked finished - for t in self.tids_marked_toretry: - finished = [ x for x in finished if x != t ] - finishedtasks = [ x for x in finishedtasks if x != t ] - - candidates = candidates + self.tids_marked_toretry - self.tids_marked_toretry = [] - - - # new candidates - for tid in finished: - if self.possiblenexttask.get(tid)!=None: - potential_candidates=list(self.possiblenexttask[tid]) - for candid in potential_candidates: - # try to see if this is really a candidate: - if self.is_good_candidate(candid, finishedtasks) and candidates.count(candid)==0: - candidates.append(candid) - - actionlogger.debug("New candidates " + str( candidates)) - send_webhook(self.args.webhook, "New candidates " + str(candidates)) - - if len(candidates)==0 and len(self.process_list)==0: - break - except Exception as e: - exc_type, exc_obj, exc_tb = sys.exc_info() - fname = os.path.split(exc_tb.tb_frame.f_code.co_filename)[1] - print(exc_type, fname, exc_tb.tb_lineno) - traceback.print_exc() - print ('Cleaning up ') - - self.SIGHandler(0,0) - - endtime = time.perf_counter() - statusmsg = "success" - if errorencountered: - statusmsg = "with failures" - - print ('\n**** Pipeline done ' + statusmsg + ' (global_runtime : {:.3f}s) *****\n'.format(endtime-self.start_time)) - actionlogger.debug("global_runtime : {:.3f}s".format(endtime-self.start_time)) - return errorencountered +_HERE = os.path.dirname(os.path.realpath(__file__)) -if args.cgroup!=None: - myPID=os.getpid() - # cgroups such as /sys/fs/cgroup/cpuset//tasks - # or /sys/fs/cgroup/cpu//tasks - command="echo " + str(myPID) + f" > {args.cgroup}" - actionlogger.info(f"Try running in cgroup {args.cgroup}") - waitstatus = os.system(command) - if code := os.waitstatus_to_exitcode(waitstatus): - actionlogger.error(f"Could not apply cgroup") - exit(code) - actionlogger.info("Running in cgroup") +RUNNERS = { + "legacy": os.path.join(_HERE, "o2dpg_workflow_runner_legacy.py"), + "new": os.path.join(_HERE, os.pardir, "workflow_runner", + "o2dpg_workflow_runner.py"), +} +DEFAULT_RUNNER = "legacy" -# This starts the fanotify fileaccess monitoring process -# if asked for -o2dpg_filegraph_exec = os.getenv("O2DPG_PRODUCE_FILEGRAPH") # switches filegraph monitoring on and contains the executable name -if o2dpg_filegraph_exec: - env = os.environ.copy() - env["FILEACCESS_MON_ROOTPATH"] = os.getcwd() - env["MAXMOTHERPID"] = f"{os.getpid()}" - fileaccess_log_file_name = f"pipeline_fileaccess_{os.getpid()}.log" - fileaccess_log_file = open(fileaccess_log_file_name, "w") - fileaccess_monitor_proc = subprocess.Popen( - [o2dpg_filegraph_exec], - stdout=fileaccess_log_file, - stderr=subprocess.STDOUT, - env=env) -else: - fileaccess_monitor_proc = None +def main(): + which = os.getenv("ALIEN_O2DPG_WORKFLOW_RUNNER", DEFAULT_RUNNER).strip() + script = RUNNERS.get(which) + if script is None: + sys.exit(f"ALIEN_O2DPG_WORKFLOW_RUNNER={which!r} is not one of " + f"{', '.join(sorted(RUNNERS))}") + os.execv(sys.executable, [sys.executable, script] + sys.argv[1:]) -try: - # This is core workflow runner invocation - executor=WorkflowExecutor(args.workflowfile,jmax=int(args.maxjobs),args=args) - rc = executor.execute() -finally: - if fileaccess_monitor_proc: - fileaccess_monitor_proc.terminate() # sends SIGTERM - try: - fileaccess_monitor_proc.wait(timeout=5) - except subprocess.TimeoutExpired: - fileaccess_monitor_proc.kill() # force kill if not stopping - # now produce the final filegraph output - o2dpg_root = os.getenv("O2DPG_ROOT") - analyse_cmd = [ - sys.executable, # runs with same Python interpreter - f"{o2dpg_root}/UTILS/FileIOGraph/analyse_FileIO.py", - "--actionFile", actionlogger_file, - "--monitorFile", fileaccess_log_file_name, - "-o", f"pipeline_fileaccess_report_{os.getpid()}.json", - "--basedir", os.getcwd() ] - print (f"Producing FileIOGraph with command {analyse_cmd}") - subprocess.run(analyse_cmd, check=True) -sys.exit(rc) \ No newline at end of file +if __name__ == "__main__": + main() diff --git a/MC/bin/o2dpg_workflow_runner_legacy.py b/MC/bin/o2dpg_workflow_runner_legacy.py new file mode 100755 index 000000000..eebde6501 --- /dev/null +++ b/MC/bin/o2dpg_workflow_runner_legacy.py @@ -0,0 +1,2004 @@ +#!/usr/bin/env python3 + +# started February 2021, sandro.wenzel@cern.ch + +import re +import subprocess +import time +import json +import logging +import os +import signal +import socket +import sys +import traceback +import platform +import tarfile +from copy import deepcopy +try: + from graphviz import Digraph + havegraphviz=True +except ImportError: + havegraphviz=False + +formatter = logging.Formatter('%(asctime)s %(levelname)s %(message)s') + +sys.setrecursionlimit(100000) + +import argparse +import psutil +max_system_mem=psutil.virtual_memory().total + +sys.path.append(os.path.join(os.path.dirname(__file__), '.', 'o2dpg_workflow_utils')) +from o2dpg_workflow_utils import read_workflow + +# defining command line options +parser = argparse.ArgumentParser(description='Parallel execution of a (O2-DPG) DAG data/job pipeline under resource contraints.', + formatter_class=argparse.ArgumentDefaultsHelpFormatter) + +parser.add_argument('-f','--workflowfile', help='Input workflow file name', required=True) +parser.add_argument('-jmax','--maxjobs', type=int, help='Number of maximal parallel tasks.', default=100) +parser.add_argument('-k','--keep-going', action='store_true', help='Keep executing the pipeline as far possibe (not stopping on first failure)') +parser.add_argument('--dry-run', action='store_true', help='Show what you would do.') +parser.add_argument('--visualize-workflow', action='store_true', help='Saves a graph visualization of workflow.') +parser.add_argument('--target-labels', nargs='+', help='Runs the pipeline by target labels (example "TPC" or "DIGI").\ + This condition is used as logical AND together with --target-tasks.', default=[]) +parser.add_argument('-tt','--target-tasks', nargs='+', help='Runs the pipeline by target tasks (example "tpcdigi"). By default everything in the graph is run. Regular expressions supported.', default=["*"]) +parser.add_argument('--produce-script', help='Produces a shell script that runs the workflow in serialized manner and quits.') +parser.add_argument('--rerun-from', help='Reruns the workflow starting from given task (or pattern). All dependent jobs will be rerun.') +parser.add_argument('--list-tasks', help='Simply list all tasks by name and quit.', action='store_true') + +# Resources +parser.add_argument('--update-resources', dest="update_resources", help='Read resource estimates from a JSON and apply where possible.') +parser.add_argument("--dynamic-resources", dest="dynamic_resources", action="store_true", help="Update reources estimates of task based on finished related tasks") # derive resources dynamically +parser.add_argument('--optimistic-resources', dest="optimistic_resources", action="store_true", help="Try to run workflow even though resource limits might underestimate resource needs of some tasks") +parser.add_argument("--n-backfill", dest="n_backfill", type=int, default=1) +parser.add_argument('--mem-limit', help='Set memory limit as scheduling constraint (in MB)', default=0.9*max_system_mem/1024./1024, type=float) +parser.add_argument('--cpu-limit', help='Set CPU limit (core count)', default=8, type=float) +parser.add_argument('--cgroup', help='Execute pipeline under a given cgroup (e.g., 8coregrid) emulating resource constraints. This m\ +ust exist and the tasks file must be writable to with the current user.') + +# run control, webhooks +parser.add_argument('--stdout-on-failure', action='store_true', help='Print log files of failing tasks to stdout,') +parser.add_argument('--webhook', help=argparse.SUPPRESS) # log some infos to this webhook channel +parser.add_argument('--checkpoint-on-failure', help=argparse.SUPPRESS) # debug option making a debug-tarball and sending to specified address + # argument is alien-path +parser.add_argument('--retry-on-failure', help=argparse.SUPPRESS, default=0) # number of times a failing task is retried +parser.add_argument('--no-rootinit-speedup', help=argparse.SUPPRESS, action='store_true') # disable init of ROOT environment vars to speedup init/startup + +parser.add_argument('--remove-files-early', type=str, default="", help="Delete intermediate files early (using the file graph information in the given file)") + + +# Logging +parser.add_argument('--action-logfile', help='Logfilename for action logs. If none given, pipeline_action_#PID.log will be used') +parser.add_argument('--metric-logfile', help='Logfilename for metric logs. If none given, pipeline_metric_#PID.log will be used') +parser.add_argument('--production-mode', action='store_true', help='Production mode') +# will trigger special features good for non-interactive/production processing (automatic cleanup of files etc). +args = parser.parse_args() + +def setup_logger(name, log_file, level=logging.INFO): + """To setup as many loggers as you want""" + + handler = logging.FileHandler(log_file, mode='w') + handler.setFormatter(formatter) + + logger = logging.getLogger(name) + logger.setLevel(level) + logger.addHandler(handler) + + return logger + +# first file logger +actionlogger_file = ('pipeline_action_' + str(os.getpid()) + '.log', args.action_logfile)[args.action_logfile!=None] +actionlogger = setup_logger('pipeline_action_logger', actionlogger_file, level=logging.DEBUG) + +# second file logger +metriclogger = setup_logger('pipeline_metric_logger', ('pipeline_metric_' + str(os.getpid()) + '.log', args.action_logfile)[args.action_logfile!=None]) + +# Immediately log imposed memory and CPU limit as well as further useful meta info +_ , meta = read_workflow(args.workflowfile) +meta["cpu_limit"] = args.cpu_limit +meta["mem_limit"] = args.mem_limit +meta["workflow_file"] = os.path.abspath(args.workflowfile) +args.target_tasks = [f.strip('"').strip("'") for f in args.target_tasks] # strip quotes from the shell +meta["target_task"] = args.target_tasks +meta["rerun_from"] = args.rerun_from +meta["target_labels"] = args.target_labels +metriclogger.info(meta) + +# for debugging without terminal access +# TODO: integrate into standard logger +def send_webhook(hook, t): + if hook!=None: + command="curl -X POST -H 'Content-type: application/json' --data '{\"text\":\" " + str(t) + "\"}' " + str(hook) + " &> /dev/null" + os.system(command) + +# A fallback solution to getting all child procs +# in case psutil has problems (PermissionError). +# It returns the same list as psutil.children(recursive=True). +def getChildProcs(basepid): + cmd=''' + childprocs() { + local parent=$1 + if [ ! "$2" ]; then + child_pid_list="" + fi + if [ "$parent" ] ; then + child_pid_list="$child_pid_list $parent" + for childpid in $(pgrep -P ${parent}); do + childprocs $childpid "nottoplevel" + done; + fi + # return via a string list (only if toplevel) + if [ ! "$2" ]; then + echo "${child_pid_list}" + fi + } + ''' + cmd = cmd + '\n' + 'childprocs ' + str(basepid) + output = subprocess.check_output(cmd, shell=True) + plist = [] + for p in output.strip().split(): + try: + proc=psutil.Process(int(p)) + except psutil.NoSuchProcess: + continue + + plist.append(proc) + return plist + +# +# Code section to find all topological orderings +# of a DAG. This is used to know when we can schedule +# things in parallel. +# Taken from https://www.geeksforgeeks.org/all-topological-sorts-of-a-directed-acyclic-graph/ + +# class to represent a graph object +class Graph: + + # Constructor + def __init__(self, edges, N): + + # A List of Lists to represent an adjacency list + self.adjList = [[] for _ in range(N)] + + # stores in-degree of a vertex + # initialize in-degree of each vertex by 0 + self.indegree = [0] * N + + # add edges to the undirected graph + for (src, dest) in edges: + + # add an edge from source to destination + self.adjList[src].append(dest) + + # increment in-degree of destination vertex by 1 + self.indegree[dest] = self.indegree[dest] + 1 + +# Recursive function to find all topological orderings of a given DAG +def findAllTopologicalOrders(graph, path, discovered, N, allpaths, maxnumber=1): + if len(allpaths) >= maxnumber: + return + + # do for every vertex + for v in range(N): + + # proceed only if in-degree of current node is 0 and + # current node is not processed yet + if graph.indegree[v] == 0 and not discovered[v]: + + # for every adjacent vertex u of v, reduce in-degree of u by 1 + for u in graph.adjList[v]: + graph.indegree[u] = graph.indegree[u] - 1 + + # include current node in the path and mark it as discovered + path.append(v) + discovered[v] = True + + # recur + findAllTopologicalOrders(graph, path, discovered, N, allpaths) + + # backtrack: reset in-degree information for the current node + for u in graph.adjList[v]: + graph.indegree[u] = graph.indegree[u] + 1 + + # backtrack: remove current node from the path and + # mark it as undiscovered + path.pop() + discovered[v] = False + + # record valid ordering + if len(path) == N: + allpaths.append(path.copy()) + + +# get all topological orderings of a given DAG as a list +def printAllTopologicalOrders(graph, maxnumber=1): + # get number of nodes in the graph + N = len(graph.adjList) + + # create an auxiliary space to keep track of whether vertex is discovered + discovered = [False] * N + + # list to store the topological order + path = [] + allpaths = [] + # find all topological ordering and print them + findAllTopologicalOrders(graph, path, discovered, N, allpaths, maxnumber=maxnumber) + return allpaths + +# <--- end code section for topological sorts + +# find all tasks that depend on a given task (id); when a cache +# dict is given we can fill for the whole graph in one pass... +def find_all_dependent_tasks(possiblenexttask, tid, cache=None): + c=cache.get(tid) if cache else None + if c!=None: + return c + + daughterlist=[tid] + # possibly recurse + for n in possiblenexttask[tid]: + c = cache.get(n) if cache else None + if c == None: + c = find_all_dependent_tasks(possiblenexttask, n, cache) + daughterlist = daughterlist + c + if cache is not None: + cache[n]=c + + if cache is not None: + cache[tid]=daughterlist + return list(set(daughterlist)) + + +# wrapper taking some edges, constructing the graph, +# obtain all topological orderings and some other helper data structures +def analyseGraph(edges, nodes): + # Number of nodes in the graph + N = len(nodes) + + # candidate list trivial + nextjobtrivial = { n:[] for n in nodes } + # startnodes + nextjobtrivial[-1] = nodes + for e in edges: + nextjobtrivial[e[0]].append(e[1]) + if nextjobtrivial[-1].count(e[1]): + nextjobtrivial[-1].remove(e[1]) + + # find topological orderings of the graph + # create a graph from edges + graph = Graph(edges, N) + orderings = printAllTopologicalOrders(graph) + + return (orderings, nextjobtrivial) + + +def draw_workflow(workflowspec): + if not havegraphviz: + print('graphviz not installed, cannot draw workflow') + return + + dot = Digraph(comment='MC workflow') + nametoindex={} + index=0 + # nodes + for node in workflowspec['stages']: + name=node['name'] + nametoindex[name]=index + dot.node(str(index), name) + index=index+1 + + # edges + for node in workflowspec['stages']: + toindex = nametoindex[node['name']] + for req in node['needs']: + fromindex = nametoindex[req] + dot.edge(str(fromindex), str(toindex)) + + dot.render('workflow.gv') + +# builds the graph given a "taskuniverse" list +# builds accompagnying structures tasktoid and idtotask +def build_graph(taskuniverse, workflowspec): + tasktoid={ t[0]['name']:i for i, t in enumerate(taskuniverse, 0) } + # print (tasktoid) + + nodes = [] + edges = [] + for t in taskuniverse: + nodes.append(tasktoid[t[0]['name']]) + for n in t[0]['needs']: + edges.append((tasktoid[n], tasktoid[t[0]['name']])) + + return (edges, nodes) + + +# loads json into dict, e.g. for workflow specification +def load_json(workflowfile): + fp=open(workflowfile) + workflowspec=json.load(fp) + return workflowspec + + +# filters the original workflowspec according to wanted targets or labels +# returns a new workflowspec and the list of "final" workflowtargets +def filter_workflow(workflowspec, targets=[], targetlabels=[]): + if len(targets)==0: + return workflowspec, [] + if len(targetlabels)==0 and len(targets)==1 and targets[0]=="*": + return workflowspec, [] + + transformedworkflowspec = workflowspec + + def task_matches(t): + for filt in targets: + if filt=="*": + return True + if re.match(filt, t) != None: + return True + return False + + def task_matches_labels(t): + # when no labels are given at all it's ok + if len(targetlabels)==0: + return True + + for l in t['labels']: + if targetlabels.count(l)!=0: + return True + return False + + # The following sequence of operations works and is somewhat structured. + # However, it builds lookups used elsewhere as well, so some CPU might be saved by reusing + # some structures across functions or by doing less passes on the data. + + # helper lookup + tasknametoid = { t['name']:i for i, t in enumerate(workflowspec['stages'],0) } + + # check if a task can be run at all + # or not due to missing requirements + def canBeDone(t,cache={}): + ok = True + c = cache.get(t['name']) + if c != None: + return c + for r in t['needs']: + taskid = tasknametoid.get(r) + if taskid != None: + if not canBeDone(workflowspec['stages'][taskid], cache): + ok = False + break + else: + ok = False + break + cache[t['name']] = ok + if ok == False: + print (f"Disabling target {t['name']} due to unsatisfied requirements") + return ok + + okcache = {} + # build full target list + full_target_list = [ t for t in workflowspec['stages'] if task_matches(t['name']) and task_matches_labels(t) and canBeDone(t,okcache) ] + full_target_name_list = [ t['name'] for t in full_target_list ] + + # build full dependency list for a task t + def getallrequirements(t): + _l=[] + for r in t['needs']: + fulltask = workflowspec['stages'][tasknametoid[r]] + _l.append(fulltask) + _l=_l+getallrequirements(fulltask) + return _l + + full_requirements_list = [ getallrequirements(t) for t in full_target_list ] + + # make flat and fetch names only + full_requirements_name_list = list(set([ item['name'] for sublist in full_requirements_list for item in sublist ])) + + # inner "lambda" helper answering if a task "name" is needed by given targets + def needed_by_targets(name): + if full_target_name_list.count(name)!=0: + return True + if full_requirements_name_list.count(name)!=0: + return True + return False + + # we finaly copy everything matching the targets as well + # as all their requirements + transformedworkflowspec['stages']=[ l for l in workflowspec['stages'] if needed_by_targets(l['name']) ] + return transformedworkflowspec, full_target_name_list + + +# builds topological orderings (for each timeframe) +def build_dag_properties(workflowspec): + globaltaskuniverse = [ (l, i) for i, l in enumerate(workflowspec['stages'], 1) ] + timeframeset = set( l['timeframe'] for l in workflowspec['stages'] ) + + edges, nodes = build_graph(globaltaskuniverse, workflowspec) + tup = analyseGraph(edges, nodes.copy()) + # + global_next_tasks = tup[1] + + + dependency_cache = {} + # weight influences scheduling order can be anything user defined ... for the moment we just prefer to stay within a timeframe + # then take the number of tasks that depend on a task as further weight + # TODO: bring in resource estimates from runtime, CPU, MEM + # TODO: make this a policy of the runner to study different strategies + def getweight(tid): + return (globaltaskuniverse[tid][0]['timeframe'], len(find_all_dependent_tasks(global_next_tasks, tid, dependency_cache))) + + task_weights = [ getweight(tid) for tid in range(len(globaltaskuniverse)) ] + + for tid in range(len(globaltaskuniverse)): + actionlogger.info("Score for " + str(globaltaskuniverse[tid][0]['name']) + " is " + str(task_weights[tid])) + + # print (global_next_tasks) + return { 'nexttasks' : global_next_tasks, 'weights' : task_weights, 'topological_ordering' : tup[0] } + + +# update the resource estimates of a workflow based on resources given via JSON +def update_resource_estimates(workflow, resource_json): + # the resource_dict here is generated by tool o2dpg_sim_metrics.py json-stat + resource_dict = load_json(resource_json) + stages = workflow["stages"] + + for task in stages: + if task["timeframe"] >= 1: + name = "_".join(task["name"].split("_")[:-1]) + else: + name = task["name"] + + if name not in resource_dict: + continue + + new_resources = resource_dict[name] + + # memory + newmem = new_resources.get("pss", {}).get("max", None) + if newmem is not None: + oldmem = task["resources"]["mem"] + actionlogger.info("Updating mem estimate for " + task["name"] + " from " + str(oldmem) + " to " + str(newmem)) + task["resources"]["mem"] = newmem + + # cpu + newcpu = new_resources.get("cpu", {}).get("mean", None) + if newcpu is not None: + oldcpu = task["resources"]["cpu"] + rel_cpu = task["resources"]["relative_cpu"] + # TODO: No longer sure about this since we inject numbers from actually measured workloads + if rel_cpu is not None: + # respect the relative CPU settings + # By default, the CPU value in the workflow is already scaled if relative_cpu is given. + # The new estimate on the other hand is not yet scaled so it needs to be done here. + newcpu *= rel_cpu + actionlogger.info("Updating cpu estimate for " + task["name"] + " from " + str(oldcpu) + " to " + str(newcpu)) + task["resources"]["cpu"] = newcpu + +# a function to read a software environment determined by alienv into +# a python dictionary +def get_alienv_software_environment(packagestring): + """ + packagestring is something like O2::v202298081-1,O2Physics::xxx representing packages + published on CVMFS ... or ... a file containing directly the software environment to apply + """ + + # the trivial cases do nothing + if packagestring == None or packagestring == "" or packagestring == "None": + return {} + + def load_env_file(env_file): + """Transform an environment file generated with 'export > env.txt' into a python dictionary.""" + env_vars = {} + with open(env_file, "r") as f: + for line in f: + line = line.strip() + + # Ignore empty lines or comments + if not line or line.startswith("#"): + continue + + # Remove 'declare -x ' if present + if line.startswith("declare -x "): + line = line.replace("declare -x ", "", 1) + + # Handle case: "FOO" without "=" (assign empty string) + if "=" not in line: + key, value = line.strip(), "" + else: + key, value = line.split("=", 1) + value = value.strip('"') # Remove surrounding quotes if present + + env_vars[key.strip()] = value + return env_vars + + # see if this is a file + if os.path.exists(packagestring) and os.path.isfile(packagestring): + actionlogger.info("Taking software environment from file " + packagestring) + return load_env_file(packagestring) + + # alienv printenv packagestring --> dictionary + # for the moment this works with CVMFS only + cmd="/cvmfs/alice.cern.ch/bin/alienv printenv " + packagestring + proc = subprocess.Popen([cmd], stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True) + + envstring, err = proc.communicate() + # see if the printenv command was successful + if len(err.decode()) > 0: + print (err.decode()) + raise Exception + + # the software environment is now in the evnstring + # split it on semicolon + envstring=envstring.decode() + tokens=envstring.split(";") + # build envmap + envmap = {} + for t in tokens: + # check if assignment + if t.count("=") > 0: + assignment = t.rstrip().split("=") + envmap[assignment[0]] = assignment[1] + elif t.count("export") > 0: + # the case when we export or a simple variable + # need to consider the case when this has not been previously assigned + variable = t.split()[1] + if not variable in envmap: + envmap[variable]="" + + return envmap + +# +# functions for execution; encapsulated in a WorkflowExecutor class +# + +class Semaphore: + """ + Object that can be used as semaphore + """ + def __init__(self): + self.locked = False + def lock(self): + self.locked = True + def unlock(self): + self.locked = False + + +class ResourceBoundaries: + """ + Container holding global resource properties + """ + def __init__(self, cpu_limit, mem_limit, dynamic_resources=False, optimistic_resources=False): + self.cpu_limit = cpu_limit + self.mem_limit = mem_limit + self.dynamic_resources = dynamic_resources + # if this is set, tasks that would normally go beyond the resource limits will tried to be run in any case + self.optimistic_resources = optimistic_resources + + +class TaskResources: + """ + Container holding resources of a single task + """ + def __init__(self, tid, name, cpu, cpu_relative, mem, resource_boundaries): + # the task ID belonging to these resources + self.tid = tid + self.name = name + # original CPUs/MEM assigned (persistent) + self.cpu_assigned_original = cpu + self.mem_assigned_original = mem + # relative CPU, to be multiplied with sampled CPU; set by the user, e.g. to allow to backfill tasks + # only takes effect when sampling resources; persistent + self.cpu_relative = cpu_relative if cpu_relative else 1 + # CPUs/MEM assigned (transient) + self.cpu_assigned = cpu + self.mem_assigned = mem + # global resource settings + self.resource_boundaries = resource_boundaries + # sampled resources of this + self.cpu_sampled = None + self.mem_sampled = None + # Set these after a task has finished to compute new estimates for related tasks + self.walltime = None + self.cpu_taken = None + self.mem_taken = None + # collected during monitoring + self.time_collect = [] + self.cpu_collect = [] + self.mem_collect = [] + # linked to other resources of task that are of the same type as this one + self.related_tasks = None + # can assign a semaphore + self.semaphore = None + # the task's nice value + self.nice_value = None + # whether or not the task's resources are currently booked + self.booked = False + + @property + def is_done(self): + return self.time_collect and not self.booked + + def is_within_limits(self): + """ + Check if assigned resources respect limits + """ + cpu_within_limits = True + mem_within_limits = True + if self.cpu_assigned > self.resource_boundaries.cpu_limit: + cpu_within_limits = False + actionlogger.warning("CPU of task %s exceeds limits %d > %d", self.name, self.cpu_assigned, self.resource_boundaries.cpu_limit) + if self.cpu_assigned > self.resource_boundaries.mem_limit: + mem_within_limits = False + actionlogger.warning("MEM of task %s exceeds limits %d > %d", self.name, self.cpu_assigned, self.resource_boundaries.cpu_limit) + return cpu_within_limits and mem_within_limits + + def limit_resources(self, cpu_limit=None, mem_limit=None): + """ + Limit resources of this specific task + """ + if not cpu_limit: + cpu_limit = self.resource_boundaries.cpu_limit + if not mem_limit: + mem_limit = self.resource_boundaries.mem_limit + self.cpu_assigned = min(self.cpu_assigned, cpu_limit) + self.mem_assigned = min(self.mem_assigned, mem_limit) + + def add(self, time_passed, cpu, mem): + """ + Brief interface to add resources that were measured after time_passed + """ + self.time_collect.append(time_passed) + self.cpu_collect.append(cpu) + self.mem_collect.append(mem) + + def sample_resources(self): + """ + If this task is done, sample CPU and MEM for all related tasks that have not started yet + """ + if not self.is_done: + return + + if len(self.time_collect) < 3: + # Consider at least 3 points to sample from + self.cpu_sampled = self.cpu_assigned + self.mem_sampled = self.mem_assigned + actionlogger.debug("Task %s has not enough points (< 3) to sample resources, setting to previosuly assigned values.", self.name) + else: + # take the time deltas and leave out the very first CPU measurent which is not meaningful, + # at least when it domes from psutil.Proc.cpu_percent(interval=None) + time_deltas = [self.time_collect[i+1] - self.time_collect[i] for i in range(len(self.time_collect) - 1)] + cpu = sum([cpu * time_delta for cpu, time_delta in zip(self.cpu_collect[1:], time_deltas) if cpu >= 0]) + self.cpu_sampled = cpu / sum(time_deltas) + self.mem_sampled = max(self.mem_collect) + + mem_sampled = 0 + cpu_sampled = [] + for res in self.related_tasks: + if res.is_done: + mem_sampled = max(mem_sampled, res.mem_sampled) + cpu_sampled.append(res.cpu_sampled) + cpu_sampled = sum(cpu_sampled) / len(cpu_sampled) + + # This task ran already with the assigned resources, so let's set it to the limit + if cpu_sampled > self.resource_boundaries.cpu_limit: + actionlogger.warning("Sampled CPU (%.2f) exceeds assigned CPU limit (%.2f)", cpu_sampled, self.resource_boundaries.cpu_limit) + elif cpu_sampled < 0: + actionlogger.debug("Sampled CPU for %s is %.2f < 0, setting to previously assigned value %.2f", self.name, cpu_sampled, self.cpu_assigned) + cpu_sampled = self.cpu_assigned + + if mem_sampled > self.resource_boundaries.mem_limit: + actionlogger.warning("Sampled MEM (%.2f) exceeds assigned MEM limit (%.2f)", mem_sampled, self.resource_boundaries.mem_limit) + elif mem_sampled <= 0: + actionlogger.debug("Sampled memory for %s is %.2f <= 0, setting to previously assigned value %.2f", self.name, mem_sampled, self.mem_assigned) + mem_sampled = self.mem_assigned + + for res in self.related_tasks: + if res.is_done or res.booked: + continue + res.cpu_assigned = cpu_sampled * res.cpu_relative + res.mem_assigned = mem_sampled + # This task has been run before, stay optimistic and limit the resources in case the sampled ones exceed limits + res.limit_resources() + + +class ResourceManager: + """ + Central class to manage resources + + - CPU limits + - MEM limits + - Semaphores + + Entrypoint to set and to query for resources to be updated. + + Can be asked whether a certain task can be run under current resource usage. + Book and unbook resources. + """ + def __init__(self, cpu_limit, mem_limit, procs_parallel_max=100, dynamic_resources=False, optimistic_resources=False): + """ + Initialise members with defaults + """ + # hold TaskResources of all tasks + self.resources = [] + + # helper dictionaries holding common objects which will be distributed to single TaskResources objects + # to avoid further lookup and at the same time to share the same common objects + self.resources_related_tasks_dict = {} + self.semaphore_dict = {} + + # one common object that holds global resource settings such as CPU and MEM limits + self.resource_boundaries = ResourceBoundaries(cpu_limit, mem_limit, dynamic_resources, optimistic_resources) + + # register resources that are booked under default nice value + self.cpu_booked = 0 + self.mem_booked = 0 + # number of tasks currently booked + self.n_procs = 0 + + # register resources that are booked under high nice value + self.cpu_booked_backfill = 0 + self.mem_booked_backfill = 0 + # number of tasks currently booked under high nice value + self.n_procs_backfill = 0 + + # the maximum number of tasks that run at the same time + self.procs_parallel_max = procs_parallel_max + + # get the default nice value of this python script + self.nice_default = os.nice(0) + # add 19 to get nice value of low-priority tasks + self.nice_backfill = self.nice_default + 19 + + def add_task_resources(self, name, related_tasks_name, cpu, cpu_relative, mem, semaphore_string=None): + """ + Construct and Add a new TaskResources object + """ + resources = TaskResources(len(self.resources), name, cpu, cpu_relative, mem, self.resource_boundaries) + if not resources.is_within_limits() and not self.resource_boundaries.optimistic_resources: + # exit if we don't dare to try + print(f"Resources of task {name} are exceeding the boundaries.\nCPU: {cpu} (estimate) vs. {self.resource_boundaries.cpu_limit} (boundary)\nMEM: {mem} (estimated) vs. {self.resource_boundaries.mem_limit} (boundary).") + print("Pass --optimistic-resources to the runner to attempt the run anyway.") + exit(1) + # if we get here, either all is good or the user decided to be optimistic and we limit the resources, by default to the given CPU and mem limits. + resources.limit_resources() + + self.resources.append(resources) + # do the following to have the same Semaphore object for all corresponding TaskResources so that we do not need a lookup + if semaphore_string: + if semaphore_string not in self.semaphore_dict: + self.semaphore_dict[semaphore_string] = Semaphore() + resources.semaphore = self.semaphore_dict[semaphore_string] + + # do the following to give each TaskResources a list of the related tasks so we do not need an additional lookup + if related_tasks_name: + if related_tasks_name not in self.resources_related_tasks_dict: + # assigned list is [valid top be used, list of CPU, list of MEM, list of walltimes of each related task, list of processes that ran in parallel on average, list of taken CPUs, list of assigned CPUs, list of tasks finished in the meantime] + self.resources_related_tasks_dict[related_tasks_name] = [] + self.resources_related_tasks_dict[related_tasks_name].append(resources) + resources.related_tasks = self.resources_related_tasks_dict[related_tasks_name] + + def add_monitored_resources(self, tid, time_delta_since_start, cpu, mem): + self.resources[tid].add(time_delta_since_start, cpu, mem) + + def book(self, tid, nice_value): + """ + Book the resources of this task with given nice value + + The final nice value is determined by the final submission and could be different. + This can happen if the nice value should have been changed while that is not allowed by the system. + """ + res = self.resources[tid] + # take the nice value that was previously assigned when resources where checked last time + previous_nice_value = res.nice_value + + if previous_nice_value is None: + # this has not been checked ever if it was ok to be submitted + actionlogger.warning("Task ID %d has never been checked for resources. Treating as backfill", tid) + nice_value = self.nice_backfill + elif res.nice_value != nice_value: + actionlogger.warning("Task ID %d has was last time checked for a different nice value (%d) but is now submitted with (%d).", tid, res.nice_value, nice_value) + + res.nice_value = nice_value + res.booked = True + if res.semaphore is not None: + res.semaphore.lock() + if nice_value != self.nice_default: + self.n_procs_backfill += 1 + self.cpu_booked_backfill += res.cpu_assigned + self.mem_booked_backfill += res.mem_assigned + return + self.n_procs += 1 + self.cpu_booked += res.cpu_assigned + self.mem_booked += res.mem_assigned + + def unbook(self, tid): + """ + Unbook the reources of this task + """ + res = self.resources[tid] + res.booked = False + if self.resource_boundaries.dynamic_resources: + res.sample_resources() + if res.semaphore is not None: + res.semaphore.unlock() + if res.nice_value != self.nice_default: + self.cpu_booked_backfill -= res.cpu_assigned + self.mem_booked_backfill -= res.mem_assigned + self.n_procs_backfill -= 1 + if self.n_procs_backfill <= 0: + self.cpu_booked_backfill = 0 + self.mem_booked_backfill = 0 + return + self.n_procs -= 1 + self.cpu_booked -= res.cpu_assigned + self.mem_booked -= res.mem_assigned + if self.n_procs <= 0: + self.cpu_booked = 0 + self.mem_booked = 0 + + def ok_to_submit(self, tids): + """ + This generator yields the tid and nice value tuple from the list of task ids that should be checked + """ + tids_copy = tids.copy() + + def ok_to_submit_default(res): + """ + Return default nice value if conditions are met, None otherwise + """ + # analyse CPU + okcpu = (self.cpu_booked + res.cpu_assigned <= self.resource_boundaries.cpu_limit) + # analyse MEM + okmem = (self.mem_booked + res.mem_assigned <= self.resource_boundaries.mem_limit) + actionlogger.debug ('Condition check --normal-- for ' + str(res.tid) + ':' + res.name + ' CPU ' + str(okcpu) + ' MEM ' + str(okmem)) + return self.nice_default if (okcpu and okmem) else None + + def ok_to_submit_backfill(res, backfill_cpu_factor=1.5, backfill_mem_factor=1.5): + """ + Return backfill nice value if conditions are met, None otherwise + """ + if self.n_procs_backfill >= args.n_backfill: + return None + + if res.cpu_assigned > 0.9 * self.resource_boundaries.cpu_limit or res.mem_assigned / self.resource_boundaries.cpu_limit >= 1900: + return None + + # analyse CPU + okcpu = (self.cpu_booked_backfill + res.cpu_assigned <= self.resource_boundaries.cpu_limit) + okcpu = okcpu and (self.cpu_booked + self.cpu_booked_backfill + res.cpu_assigned <= backfill_cpu_factor * self.resource_boundaries.cpu_limit) + # analyse MEM + okmem = (self.mem_booked + self.mem_booked_backfill + res.mem_assigned <= backfill_mem_factor * self.resource_boundaries.mem_limit) + actionlogger.debug ('Condition check --backfill-- for ' + str(res.tid) + ':' + res.name + ' CPU ' + str(okcpu) + ' MEM ' + str(okmem)) + + return self.nice_backfill if (okcpu and okmem) else None + + if self.n_procs + self.n_procs_backfill >= self.procs_parallel_max: + # in this case, nothing can be done + return + + for ok_to_submit_impl, should_break in ((ok_to_submit_default, True), (ok_to_submit_backfill, False)): + tid_index = 0 + while tid_index < len(tids_copy): + + tid = tids_copy[tid_index] + res = self.resources[tid] + + actionlogger.info("Setup resources for task %s, cpu: %f, mem: %f", res.name, res.cpu_assigned, res.mem_assigned) + tid_index += 1 + + if (res.semaphore is not None and res.semaphore.locked) or res.booked: + continue + + nice_value = ok_to_submit_impl(res) + if nice_value is not None: + # if we get a non-None nice value, it means that this task is good to go + res.nice_value = nice_value + # yield the tid and its assigned nice value + yield tid, nice_value + + elif should_break: + # break here if resources of the next task do not fit + break + + +def filegraph_expand_timeframes(data: dict, timeframes: set, target_namelist) -> dict: + """ + A utility function for the fileaccess logic. Takes a template and duplicates + for the multi-timeframe structure. + """ + tf_entries = [ + entry for entry in data.get("file_report", []) + if re.match(r"^\./tf\d+/", entry["file"]) + ] + + result = {} + for i in timeframes: + if i == -1: + continue + # Deepcopy to avoid modifying original + new_entries = deepcopy(tf_entries) + for entry in new_entries: + # Fix filepath + entry["file"] = re.sub(r"^\./tf\d+/", f"./tf{i}/", entry["file"]) + # Fix written_by and read_by (preserve prefix, change numeric suffix) + entry["written_by"] = [ + re.sub(r"_\d+$", f"_{i}", w) for w in entry["written_by"] + ] + # for now we mark some files as keep if they are written + # by a target in the runner targetlist. TODO: Add other mechanisms + # to ask for file keeping (such as via regex or the like) + for e in entry["written_by"]: + if e in target_namelist: + entry["keep"] = True + entry["read_by"] = [ + re.sub(r"_\d+$", f"_{i}", r) for r in entry["read_by"] + ] + result[f"timeframe-{i}"] = new_entries + + return result + + + +class WorkflowExecutor: + # Constructor + def __init__(self, workflowfile, args, jmax=100): + self.args=args + self.is_productionmode = args.production_mode == True # os.getenv("ALIEN_PROC_ID") != None + self.workflowfile = workflowfile + self.workflowspec = load_json(workflowfile) + self.globalinit = self.extract_global_environment(self.workflowspec) # initialize global environment settings + for e in self.globalinit['env']: + if os.environ.get(e, None) == None: + value = self.globalinit['env'][e] + actionlogger.info("Applying global environment from init section " + str(e) + " : " + str(value)) + os.environ[e] = str(value) + + # only keep those tasks that are necessary to be executed based on user's filters + self.full_target_namelist = [] + self.workflowspec, self.full_target_namelist = filter_workflow(self.workflowspec, args.target_tasks, args.target_labels) + + if not self.workflowspec['stages']: + if args.target_tasks: + print ('Apparently some of the chosen target tasks are not in the workflow') + exit (0) + print ('Workflow is empty. Nothing to do') + exit (0) + + # construct the DAG, compute task weights + workflow = build_dag_properties(self.workflowspec) + if args.visualize_workflow: + draw_workflow(self.workflowspec) + self.possiblenexttask = workflow['nexttasks'] + self.taskweights = workflow['weights'] + self.topological_orderings = workflow['topological_ordering'] + self.taskuniverse = [ l['name'] for l in self.workflowspec['stages'] ] + # construct task ID <-> task name lookup + self.idtotask = [ 0 for _ in self.taskuniverse ] + self.tasktoid = {} + self.idtotf = [ l['timeframe'] for l in self.workflowspec['stages'] ] + for i, name in enumerate(self.taskuniverse): + self.tasktoid[name]=i + self.idtotask[i]=name + + if args.update_resources: + update_resource_estimates(self.workflowspec, args.update_resources) + + # construct the object that is in charge of resource management... + self.resource_manager = ResourceManager(args.cpu_limit, args.mem_limit, args.maxjobs, args.dynamic_resources, args.optimistic_resources) + for task in self.workflowspec['stages']: + # ...and add all initial resource estimates + global_task_name = self.get_global_task_name(task["name"]) + try: + cpu_relative = float(task["resources"]["relative_cpu"]) + except TypeError: + cpu_relative = 1 + self.resource_manager.add_task_resources(task["name"], global_task_name, float(task["resources"]["cpu"]), cpu_relative, float(task["resources"]["mem"]), task.get("semaphore")) + + self.procstatus = { tid:'ToDo' for tid in range(len(self.workflowspec['stages'])) } + self.taskneeds= { t:set(self.getallrequirements(t)) for t in self.taskuniverse } + self.stoponfailure = not args.keep_going + print ("Stop on failure ",self.stoponfailure) + + self.scheduling_iteration = 0 # count how often it was tried to schedule new tasks + self.process_list = [] # list of currently scheduled tasks with normal priority + self.backfill_process_list = [] # list of curently scheduled tasks with low backfill priority (not sure this is needed) + self.pid_to_psutilsproc = {} # cache of putilsproc for resource monitoring + self.pid_to_files = {} # we can auto-detect what files are produced by which task (at least to some extent) + self.pid_to_connections = {} # we can auto-detect what connections are opened by which task (at least to some extent) + signal.signal(signal.SIGINT, self.SIGHandler) + signal.siginterrupt(signal.SIGINT, False) + self.internalmonitorcounter = 0 # internal use + self.internalmonitorid = 0 # internal use + self.tids_marked_toretry = [] # sometimes we might want to retry a failed task (simply because it was "unlucky") and we put them here + self.retry_counter = [ 0 for tid in range(len(self.taskuniverse)) ] # we keep track of many times retried already + self.task_retries = [ self.workflowspec['stages'][tid].get('retry_count',0) for tid in range(len(self.taskuniverse)) ] # the per task specific "retry" number -> needs to be parsed from the JSON + + self.alternative_envs = {} # mapping of taskid to alternative software envs (to be applied on a per-task level) + # init alternative software environments + self.init_alternative_software_environments() + + # initialize container to keep track of file-task relationsships + self.file_removal_candidates = {} + self.do_early_file_removal = False + self.timeframeset = set([ task["timeframe"] for task in self.workflowspec['stages'] ]) + if args.remove_files_early != "": + with open(args.remove_files_early) as f: + filegraph_data = json.load(f) + self.do_early_file_removal = True + self.file_removal_candidates = filegraph_expand_timeframes(filegraph_data, self.timeframeset, self.full_target_namelist) + + def apply_global_env(self, environ_dict): + for e in self.globalinit['env']: + if environ_dict.get(e, None) == None: + value = self.globalinit['env'][e] + actionlogger.info("Applying global environment from init section " + str(e) + " : " + str(value)) + environ_dict[e] = str(value) + + def perform_early_file_removal(self, taskids): + """ + This function checks which files can be deleted upon completion of task + and optionally does so. + """ + + def remove_if_exists(filepath: str) -> None: + """ + Check if a file exists, and remove it if found. + """ + if os.path.exists(filepath): + fsize = os.path.getsize(filepath) + os.remove(filepath) + actionlogger.info(f"Removing {filepath} since no longer needed. Freeing {fsize/1024./1024.} MB.") + return True + + return False + + def remove_for_task_id(taskname, file_dict, timeframe_id, listofalltimeframes): + marked_for_removal = [] + + timeframestoscan = [ timeframe_id ] + if timeframe_id == -1: + timeframestoscan = [ i for i in listofalltimeframes if i != -1 ] + + # TODO: Note that this traversal of files is not certainly not optimal + # We should (and will) keep an mapping of tasks->potential files and just + # scan these. This is already provided by the FileIOGraph analysis tool. + for tid in timeframestoscan: + for i,file_entry in enumerate(file_dict[f"timeframe-{tid}"]): + filename = file_entry['file'] + read_by = file_entry['read_by'] + written_by = file_entry['written_by'] + if taskname in read_by: + file_entry['read_by'].remove(taskname) + if taskname in written_by: + file_entry['written_by'].remove(taskname) + + # TODO: in principle the written_by criterion might not be needed + if len(file_entry['read_by']) == 0 and len(file_entry['written_by']) == 0 and file_entry.get('keep', False) == False: + # the filename mentioned here is no longer needed and we can remove it + # make sure it is there and then delete it + if remove_if_exists(filename): + # also take out the file entry from the dict altogether + marked_for_removal.append(file_entry) + + #for k in marked_for_removal: + # file_dict[f"timeframe-{tid}"].remove(k) + + for tid in taskids: + taskname = self.idtotask[tid] + timeframe_id = self.idtotf[tid] + remove_for_task_id(taskname, self.file_removal_candidates, timeframe_id, self.timeframeset) + + + def SIGHandler(self, signum, frame): + """ + basically forcing shut down of all child processes + """ + actionlogger.info("Signal " + str(signum) + " caught") + try: + procs = psutil.Process().children(recursive=True) + except (psutil.NoSuchProcess): + pass + except (psutil.AccessDenied, PermissionError): + procs = getChildProcs(os.getpid()) + + for p in procs: + actionlogger.info("Terminating " + str(p)) + try: + p.terminate() + except (psutil.NoSuchProcess, psutil.AccessDenied): + pass + + _, alive = psutil.wait_procs(procs, timeout=3) + for p in alive: + try: + actionlogger.info("Killing " + str(p)) + p.kill() + except (psutil.NoSuchProcess, psutil.AccessDenied): + pass + + exit (1) + + def extract_global_environment(self, workflowspec): + """ + Checks if the workflow contains a dedicated init task + defining a global environment. Extract information and remove from workflowspec. + """ + init_index = 0 # this has to be the first task in the workflow + globalenv = {} + initcmd = None + if workflowspec['stages'][init_index]['name'] == '__global_init_task__': + env = workflowspec['stages'][init_index].get('env', None) + if env != None: + globalenv = { e : env[e] for e in env } + cmd = workflowspec['stages'][init_index].get('cmd', None) + if cmd != 'NO-COMMAND': + initcmd = cmd + + del workflowspec['stages'][init_index] + + return {"env" : globalenv, "cmd" : initcmd } + + def execute_globalinit_cmd(self, cmd): + actionlogger.info("Executing global setup cmd " + str(cmd)) + # perform the global init command (think of cleanup/setup things to be done in any case) + p = subprocess.Popen(['/bin/bash','-c', cmd], stdout=subprocess.PIPE, stderr=subprocess.PIPE) + stdout, stderr = p.communicate() + + # Check if the command was successful (return code 0) + if p.returncode == 0: + actionlogger.info(stdout.decode()) + else: + # this should be an error + actionlogger.error("Error executing global init function") + return False + return True + + def get_global_task_name(self, name): + """ + Get the global task name + + Tasks are related if only the suffix _ is different + """ + tokens = name.split("_") + try: + int(tokens[-1]) + return "_".join(tokens[:-1]) + except ValueError: + pass + return name + + def getallrequirements(self, task_name): + """ + get all requirement of a task by its name + """ + l=[] + for required_task_name in self.workflowspec['stages'][self.tasktoid[task_name]]['needs']: + l.append(required_task_name) + l=l+self.getallrequirements(required_task_name) + return l + + def get_logfile(self, tid): + """ + O2 taskwrapper logs task stdout and stderr to logfile .log + Get its exact path based on task ID + """ + # determines the logfile name for this task + name = self.workflowspec['stages'][tid]['name'] + workdir = self.workflowspec['stages'][tid]['cwd'] + return os.path.join(workdir, f"{name}.log") + + def get_done_filename(self, tid): + """ + O2 taskwrapper leaves .log_done after a task has successfully finished + Get its exact path based on task ID + """ + return f"{self.get_logfile(tid)}_done" + + def get_resources_filename(self, tid): + """ + O2 taskwrapper leaves .log_time after a task is done + Get its exact path based on task ID + """ + return f"{self.get_logfile(tid)}_time" + + # removes the done flag from tasks that need to be run again + def remove_done_flag(self, listoftaskids): + """ + Remove .log_done files to given task IDs + """ + for tid in listoftaskids: + done_filename = self.get_done_filename(tid) + name=self.workflowspec['stages'][tid]['name'] + if args.dry_run: + print ("Would mark task " + name + " as to be done again") + else: + print ("Marking task " + name + " as to be done again") + if os.path.exists(done_filename) and os.path.isfile(done_filename): + os.remove(done_filename) + + # submits a task as subprocess and records Popen instance + def submit(self, tid, nice): + """ + Submit a task + + 1. if needed, construct working directory if it does not yet exist + 2. update lookup structures flagging the task as being run + 3. set specific environment if requested for task + 4. construct psutil.Process from command line + 4.1 adjust the niceness of that process if requested + 5. return psutil.Process object + """ + actionlogger.debug("Submitting task " + str(self.idtotask[tid]) + " with nice value " + str(nice)) + c = self.workflowspec['stages'][tid]['cmd'] + workdir = self.workflowspec['stages'][tid]['cwd'] + if workdir: + if os.path.exists(workdir) and not os.path.isdir(workdir): + actionlogger.error('Cannot create working dir ... some other resource exists already') + return None + + if not os.path.isdir(workdir): + os.makedirs(workdir) + + self.procstatus[tid]='Running' + if args.dry_run: + drycommand="echo \' " + str(self.scheduling_iteration) + " : would do " + str(self.workflowspec['stages'][tid]['name']) + "\'" + return psutil.Popen(['/bin/bash','-c',drycommand], cwd=workdir) + + taskenv = os.environ.copy() + # apply specific (non-default) software version, if any + # (this was setup earlier) + alternative_env = self.alternative_envs.get(tid, None) + if alternative_env != None and len(alternative_env) > 0: + actionlogger.info('Applying alternative software environment to task ' + self.idtotask[tid]) + if alternative_env.get('TERM') != None: + # the environment is a complete environment + taskenv = {} + taskenv = alternative_env + else: + for entry in alternative_env: + # overwrite what is present in default + taskenv[entry] = alternative_env[entry] + + # add task specific environment + if self.workflowspec['stages'][tid].get('env')!=None: + taskenv.update(self.workflowspec['stages'][tid]['env']) + + # add global workflow environment + self.apply_global_env(taskenv) + + if os.environ.get('PIPELINE_RUNNER_DUMP_TASKENVS') != None: + envfilename = "taskenv_" + str(tid) + ".log" + with open(envfilename, "w") as file: + json.dump(taskenv, file, indent=2) + + p = psutil.Popen(['/bin/bash','-c',c], cwd=workdir, env=taskenv) + try: + p.nice(nice) + except (psutil.NoSuchProcess, psutil.AccessDenied): + actionlogger.error('Couldn\'t set nice value of ' + str(p.pid) + ' to ' + str(nice)) + + return p + + def ok_to_skip(self, tid): + """ + Decide if task can be skipped based on existence of .log_done + """ + done_filename = self.get_done_filename(tid) + if os.path.exists(done_filename) and os.path.isfile(done_filename): + return True + return False + + def try_job_from_candidates(self, taskcandidates, finished): + """ + Try to schedule next tasks + + Args: + taskcandidates: list + list of possible tasks that can be submitted + finished: list + empty list that will be filled with IDs of tasks that were finished in the meantime + """ + self.scheduling_iteration = self.scheduling_iteration + 1 + + # remove "done / skippable" tasks immediately + for tid in taskcandidates.copy(): # <--- the copy is important !! otherwise this loop is not doing what you think + if self.ok_to_skip(tid): + finished.append(tid) + taskcandidates.remove(tid) + actionlogger.info("Skipping task " + str(self.idtotask[tid])) + + # if tasks_skipped: + # return # ---> we return early in order to preserve some ordering (the next candidate tried should be daughters of skipped jobs) + # get task ID and proposed niceness from generator + for (tid, nice_value) in self.resource_manager.ok_to_submit(taskcandidates): + actionlogger.debug ("trying to submit " + str(tid) + ':' + str(self.idtotask[tid])) + if p := self.submit(tid, nice_value): + # explicitly set the nice value here from the process again because it might happen that submit could not change the niceness + # so we let the ResourceManager know what the final niceness is + self.resource_manager.book(tid, p.nice()) + self.process_list.append((tid,p)) + taskcandidates.remove(tid) + # minimal delay + time.sleep(0.1) + + def stop_pipeline_and_exit(self, process_list): + # kill all remaining jobs + for p in process_list: + p[1].kill() + + exit(1) + + + def monitor(self, process_list): + """ + Go through all running tasks and get their current resources + + Resources are summed up for tasks and all their children + + Pass CPU, PSS, USS, niceness, current time to metriclogger + + Warn if overall PSS exceeds assigned memory limit + """ + self.internalmonitorcounter+=1 + if self.internalmonitorcounter % 5 != 0: + return + + self.internalmonitorid+=1 + + globalCPU=0. + globalPSS=0. + resources_per_task = {} + + # On a global level, we are interested in total disc space used (not differential in tasks) + # We can call system "du" as the fastest impl + def disk_usage_du(path: str) -> int: + """Use system du to get total size in bytes.""" + out = subprocess.check_output(['du', '-sb', path], text=True) + return int(out.split()[0]) + + disc_usage = -1 + if os.getenv("MONITOR_DISC_USAGE"): + disc_usage = disk_usage_du(os.getcwd()) / 1024. / 1024 # in MB + + for tid, proc in process_list: + + # proc is Popen object + pid=proc.pid + if self.pid_to_files.get(pid)==None: + self.pid_to_files[pid]=set() + self.pid_to_connections[pid]=set() + try: + psutilProcs = [ proc ] + # use psutil for CPU measurement + psutilProcs = psutilProcs + proc.children(recursive=True) + except (psutil.NoSuchProcess): + continue + + except (psutil.AccessDenied, PermissionError): + psutilProcs = psutilProcs + getChildProcs(pid) + + # accumulate total metrics (CPU, memory) + totalCPU = 0. + totalPSS = 0. + totalSWAP = 0. + totalUSS = 0. + for p in psutilProcs: + """ + try: + for f in p.open_files(): + self.pid_to_files[pid].add(str(f.path)+'_'+str(f.mode)) + for f in p.connections(kind="all"): + remote=f.raddr + if remote==None: + remote='none' + self.pid_to_connections[pid].add(str(f.type)+"_"+str(f.laddr)+"_"+str(remote)) + except Exception: + pass + """ + thispss=0 + thisuss=0 + # MEMORY part + try: + fullmem=p.memory_full_info() + thispss=getattr(fullmem,'pss',0) #<-- pss not available on MacOS + totalPSS=totalPSS + thispss + totalSWAP=totalSWAP + fullmem.swap + thisuss=fullmem.uss + totalUSS=totalUSS + thisuss + except (psutil.NoSuchProcess, psutil.AccessDenied): + pass + + # CPU part + # fetch existing proc or insert + cachedproc = self.pid_to_psutilsproc.get(p.pid) + if cachedproc!=None: + try: + thiscpu = cachedproc.cpu_percent(interval=None) + except (psutil.NoSuchProcess, psutil.AccessDenied): + thiscpu = 0. + totalCPU = totalCPU + thiscpu + # thisresource = {'iter':self.internalmonitorid, 'pid': p.pid, 'cpu':thiscpu, 'uss':thisuss/1024./1024., 'pss':thispss/1024./1024.} + # metriclogger.info(thisresource) + else: + self.pid_to_psutilsproc[p.pid] = p + try: + self.pid_to_psutilsproc[p.pid].cpu_percent() + except (psutil.NoSuchProcess, psutil.AccessDenied): + pass + + time_delta = int((time.perf_counter() - self.start_time) * 1000) + totalUSS = totalUSS / 1024 / 1024 + totalPSS = totalPSS / 1024 / 1024 + nice_value = proc.nice() + resources_per_task[tid]={'iter':self.internalmonitorid, + 'name':self.idtotask[tid], + 'cpu':totalCPU, + 'uss':totalUSS, + 'pss':totalPSS, + 'nice':nice_value, + 'swap':totalSWAP, + 'label':self.workflowspec['stages'][tid]['labels'], + 'disc': disc_usage} + self.resource_manager.add_monitored_resources(tid, time_delta, totalCPU / 100, totalPSS) + if nice_value == self.resource_manager.nice_default: + globalCPU += totalCPU + globalPSS += totalPSS + + metriclogger.info(resources_per_task[tid]) + send_webhook(self.args.webhook, resources_per_task) + + if globalPSS > self.resource_manager.resource_boundaries.mem_limit: + metriclogger.info('*** MEMORY LIMIT PASSED !! ***') + # --> We could use this for corrective actions such as killing jobs currently back-filling + # (or better hibernating) + + def waitforany(self, process_list, finished, failingtasks): + """ + Loop through all submitted tasks and check if they are finished + + 1. If process is still running, do nothing + 2. If process is finished, get its return value, update finished and failingtasks lists + 2.1 unbook resources + 2.2 add taken resources and pass the to ResourceManager + """ + failuredetected = False + failingpids = [] + if len(process_list)==0: + return False + + for p in list(process_list): + pid = p[1].pid + tid = p[0] # the task id of this process + returncode = 0 + if not self.args.dry_run: + returncode = p[1].poll() + if returncode!=None: + actionlogger.info ('Task ' + str(pid) + ' ' + str(tid)+':'+str(self.idtotask[tid]) + ' finished with status ' + str(returncode)) + # account for cleared resources + self.resource_manager.unbook(tid) + self.procstatus[tid]='Done' + finished.append(tid) + #self.validate_resources_running(tid) + process_list.remove(p) + if returncode != 0: + print (str(self.idtotask[tid]) + ' failed ... checking retry') + # we inspect if this is something "unlucky" which could be resolved by a simple resubmit + if self.is_worth_retrying(tid) and ((self.retry_counter[tid] < int(args.retry_on_failure)) or (self.retry_counter[tid] < int(self.task_retries[tid]))): + print (str(self.idtotask[tid]) + ' to be retried') + actionlogger.info ('Task ' + str(self.idtotask[tid]) + ' failed but marked to be retried ') + self.tids_marked_toretry.append(tid) + self.retry_counter[tid] += 1 + + else: + failuredetected = True + failingpids.append(pid) + failingtasks.append(tid) + + if failuredetected and self.stoponfailure: + actionlogger.info('Stoping pipeline due to failure in stages with PID ' + str(failingpids)) + # self.analyse_files_and_connections() + if self.args.stdout_on_failure: + self.cat_logfiles_tostdout(failingtasks) + self.send_checkpoint(failingtasks, self.args.checkpoint_on_failure) + self.stop_pipeline_and_exit(process_list) + + # empty finished means we have to wait more + return len(finished)==0 + + def is_worth_retrying(self, tid): + # This checks for some signatures in logfiles that indicate that a retry of this task + # might have a chance. + # Ideally, this should be made user configurable. Either the user could inject a lambda + # or a regular expression to use. For now we just put a hard coded list + logfile = self.get_logfile(tid) + + return True #! --> for now we just retry tasks a few times + + # 1) ZMQ_EVENT + interrupted system calls (DPL bug during shutdown) + # Not sure if grep is faster than native Python text search ... + # status = os.system('grep "failed setting ZMQ_EVENTS" ' + logfile + ' &> /dev/null') + # if os.WEXITSTATUS(status) == 0: + # return True + + # return False + + + def cat_logfiles_tostdout(self, taskids): + # In case of errors we can cat the logfiles for this taskname + # to stdout. Assuming convention that "taskname" translates to "taskname.log" logfile. + for tid in taskids: + logfile = self.get_logfile(tid) + if os.path.exists(logfile): + print (' ----> START OF LOGFILE ', logfile, ' -----') + os.system('cat ' + logfile) + print (' <---- END OF LOGFILE ', logfile, ' -----') + + def send_checkpoint(self, taskids, location): + # Makes a tarball containing all files in the base dir + # (timeframe independent) and the dir with corrupted timeframes + # and copies it to a specific ALIEN location. Not a core function + # just some tool get hold on error conditions appearing on the GRID. + + def get_tar_command(dir='./', flags='cf', findtype='f', filename='checkpoint.tar'): + return 'find ' + str(dir) + ' -maxdepth 1 -type ' + str(findtype) + ' -print0 | xargs -0 tar ' + str(flags) + ' ' + str(filename) + + if location != None: + print ('Making a failure checkpoint') + # let's determine a filename from ALIEN_PROC_ID - hostname - and PID + + aliprocid=os.environ.get('ALIEN_PROC_ID') + if aliprocid == None: + aliprocid = 0 + + fn='pipeline_checkpoint_ALIENPROC' + str(aliprocid) + '_PID' + str(os.getpid()) + '_HOST' + socket.gethostname() + '.tar' + actionlogger.info("Checkpointing to file " + fn) + tarcommand = get_tar_command(filename=fn) + actionlogger.info("Taring " + tarcommand) + + # create a README file with instruction on how to use checkpoint + readmefile=open('README_CHECKPOINT_PID' + str(os.getpid()) + '.txt','w') + + for tid in taskids: + taskspec = self.workflowspec['stages'][tid] + name = taskspec['name'] + readmefile.write('Checkpoint created because of failure in task ' + name + '\n') + readmefile.write('In order to reproduce with this checkpoint, do the following steps:\n') + readmefile.write('a) setup the appropriate O2sim environment using alienv\n') + readmefile.write('b) run: $O2DPG_ROOT/MC/bin/o2_dpg_workflow_runner.py -f workflow.json -tt ' + name + '$ --retry-on-failure 0\n') + readmefile.close() + + # first of all the base directory + os.system(tarcommand) + + # then we add stuff for the specific timeframes ids if any + for tid in taskids: + taskspec = self.workflowspec['stages'][tid] + directory = taskspec['cwd'] + if directory != "./": + tarcommand = get_tar_command(dir=directory, flags='rf', filename=fn) + actionlogger.info("Tar command is " + tarcommand) + os.system(tarcommand) + # same for soft links + tarcommand = get_tar_command(dir=directory, flags='rf', findtype='l', filename=fn) + actionlogger.info("Tar command is " + tarcommand) + os.system(tarcommand) + + # prepend file:/// to denote local file + fn = "file://" + fn + actionlogger.info("Local checkpoint file is " + fn) + + # location needs to be an alien path of the form alien:///foo/bar/ + copycommand='alien.py cp ' + fn + ' ' + str(location) + '@disk:1' + actionlogger.info("Copying to alien " + copycommand) + os.system(copycommand) + + def init_alternative_software_environments(self): + """ + Initialises alternative software environments for specific tasks, if there + is an annotation in the workflow specificiation. + """ + + environment_cache = {} + # go through all the tasks once and setup environment + for taskid in range(len(self.workflowspec['stages'])): + packagestr = self.workflowspec['stages'][taskid].get("alternative_alienv_package") + if packagestr == None: + continue + + if environment_cache.get(packagestr) == None: + environment_cache[packagestr] = get_alienv_software_environment(packagestr) + + self.alternative_envs[taskid] = environment_cache[packagestr] + + + def analyse_files_and_connections(self): + for p,s in self.pid_to_files.items(): + for f in s: + print("F" + str(f) + " : " + str(p)) + for p,s in self.pid_to_connections.items(): + for c in s: + print("C" + str(c) + " : " + str(p)) + #print(str(p) + " CONS " + str(c)) + try: + # check for intersections + for p1, s1 in self.pid_to_files.items(): + for p2, s2 in self.pid_to_files.items(): + if p1!=p2: + if type(s1) is set and type(s2) is set: + if len(s1)>0 and len(s2)>0: + try: + inters = s1.intersection(s2) + except Exception: + print ('Exception during intersect inner') + pass + if (len(inters)>0): + print ('FILE Intersection ' + str(p1) + ' ' + str(p2) + ' ' + str(inters)) + # check for intersections + for p1, s1 in self.pid_to_connections.items(): + for p2, s2 in self.pid_to_connections.items(): + if p1!=p2: + if type(s1) is set and type(s2) is set: + if len(s1)>0 and len(s2)>0: + try: + inters = s1.intersection(s2) + except Exception: + print ('Exception during intersect inner') + pass + if (len(inters)>0): + print ('CON Intersection ' + str(p1) + ' ' + str(p2) + ' ' + str(inters)) + + # check for intersections + #for p1, s1 in slf.pid_to_files.items(): + # for p2, s2 in self.pid_to_files.items(): + # if p1!=p2 and len(s1.intersection(s2))!=0: + # print ('Intersection found files ' + str(p1) + ' ' + str(p2) + ' ' + s1.intersection(s2)) + except Exception as e: + exc_type, exc_obj, exc_tb = sys.exc_info() + fname = os.path.split(exc_tb.tb_frame.f_code.co_filename)[1] + print(exc_type, fname, exc_tb.tb_lineno) + print('Exception during intersect outer') + pass + + def is_good_candidate(self, candid, finishedtasks): + if self.procstatus[candid] != 'ToDo': + return False + needs = set([self.tasktoid[t] for t in self.taskneeds[self.idtotask[candid]]]) + if set(finishedtasks).intersection(needs) == needs: + return True + return False + + def emit_code_for_task(self, tid, lines): + actionlogger.debug("Submitting task " + str(self.idtotask[tid])) + taskspec = self.workflowspec['stages'][tid] + c = taskspec['cmd'] + workdir = taskspec['cwd'] + env = taskspec.get('env') + # in general: + # try to make folder + lines.append('[ ! -d ' + workdir + ' ] && mkdir ' + workdir + '\n') + # cd folder + lines.append('cd ' + workdir + '\n') + # set local environment + if env!=None: + for e in env.items(): + lines.append('export ' + e[0] + '=' + str(e[1]) + '\n') + # do command + lines.append(c + '\n') + # unset local environment + if env!=None: + for e in env.items(): + lines.append('unset ' + e[0] + '\n') + + # cd back + lines.append('cd $OLDPWD\n') + + + # produce a bash script that runs workflow standalone + def produce_script(self, filename): + # pick one of the correct task orderings + taskorder = self.topological_orderings[0] + outF = open(filename, "w") + + lines=[] + # header + lines.append('#!/usr/bin/env bash\n') + lines.append('#THIS FILE IS AUTOGENERATED\n') + lines.append('export JOBUTILS_SKIPDONE=ON\n') + + # we record the global environment setting + # in particular to capture global workflow initialization + lines.append('#-- GLOBAL INIT SECTION FROM WORKFLOW --\n') + for e in self.globalinit['env']: + lines.append('export ' + str(e) + '=' + str(self.globalinit['env'][e]) + '\n') + lines.append('#-- TASKS FROM WORKFLOW --\n') + for tid in taskorder: + print ('Doing task ' + self.idtotask[tid]) + self.emit_code_for_task(tid, lines) + + outF.writelines(lines) + outF.close() + + def production_endoftask_hook(self, tid): + # Executes a hook at end of a successful task, meant to be used in GRID productions. + # For the moment, archiving away log files, done + time files from jobutils. + # TODO: In future this may be much more generic tasks such as dynamic cleanup of intermediate + # files (when they are no longer needed). + # TODO: Care must be taken with the continue feature as `_done` files are stored elsewhere now + actionlogger.info("Cleaning up log files for task " + str(tid)) + logf = self.get_logfile(tid) + donef = self.get_done_filename(tid) + timef = logf + "_time" + + # add to tar file archive + tf = tarfile.open(name="pipeline_log_archive.log.tar", mode='a') + if tf != None: + tf.add(logf) + tf.add(donef) + tf.add(timef) + tf.close() + + # remove original file + os.remove(logf) + os.remove(donef) + os.remove(timef) + + # print error message when no progress can be made + def noprogress_errormsg(self): + # TODO: rather than writing this out here; refer to the documentation discussion this? + msg = """Scheduler runtime error: The scheduler is not able to make progress although we have a non-zero candidate set. + +Explanation: This is typically the case because the **ESTIMATED** resource requirements for some tasks +in the workflow exceed the available number of CPU cores or the memory (as explicitely or implicitely determined from the +--cpu-limit and --mem-limit options). Often, this might be the case on laptops with <=16GB of RAM if one of the tasks +is demanding ~16GB. In this case, one could try to tell the scheduler to use a slightly higher memory limit +with an explicit --mem-limit option (for instance `--mem-limit 20000` to set to 20GB). This might work whenever the +**ACTUAL** resource usage of the tasks is smaller than anticipated (because only small test cases are run). + +In addition it might be worthwile running the workflow without this resource aware, dynamic scheduler. +This is possible by converting the json workflow into a linearized shell script and by directly executing the shell script. +Use the `--produce-script myscript.sh` option for this. +""" + print (msg, file=sys.stderr) + + def execute(self): + self.start_time = time.perf_counter() + psutil.cpu_percent(interval=None) + os.environ['JOBUTILS_SKIPDONE'] = "ON" + errorencountered = False + + def speedup_ROOT_Init(): + """initialize some env variables that speed up ROOT init + and prevent ROOT from spawning many short-lived child + processes""" + + # only do it on Linux + if platform.system() != 'Linux': + return + + if os.environ.get('ROOT_LDSYSPATH')!=None and os.environ.get('ROOT_CPPSYSINCL')!=None: + # do nothing if already defined + return + + # a) the PATH for system libraries + # search taken from ROOT TUnixSystem + cmd='LD_DEBUG=libs LD_PRELOAD=DOESNOTEXIST ls /tmp/DOESNOTEXIST 2>&1 | grep -m 1 "system search path" | sed \'s/.*=//g\' | awk \'//{print $1}\'' + proc = subprocess.Popen([cmd], stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True) + libpath, err = proc.communicate() + if not (args.no_rootinit_speedup == True): + print ("setting up ROOT system") + os.environ['ROOT_LDSYSPATH'] = libpath.decode() + os.environ['CLING_LDSYSPATH'] = libpath.decode() + + # b) the PATH for compiler includes needed by Cling + cmd = "LC_ALL=C c++ -xc++ -E -v /dev/null 2>&1 | sed -n '/^#include/,${/^ \\/.*++/{p}}'" + proc = subprocess.Popen([cmd], stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True) + incpath, err = proc.communicate() + incpaths = [ line.lstrip() for line in incpath.decode().splitlines() ] + joined = ':'.join(incpaths) + if not (args.no_rootinit_speedup == True): + actionlogger.info("Determined ROOT_CPPSYSINCL=" + joined) + os.environ['ROOT_CPPSYSINCL'] = joined + os.environ['CLING_CPPSYSINCL'] = joined + + speedup_ROOT_Init() + + # we make our own "tmp" folder + # where we can put stuff such as tmp socket files etc (for instance DPL FAIR-MQ sockets) + # (In case of running within docker/singularity, this may not be so important) + if not os.path.isdir("./.tmp"): + os.mkdir("./.tmp") + if os.environ.get('FAIRMQ_IPC_PREFIX')==None: + socketpath = os.getcwd() + "/.tmp" + actionlogger.info("Setting FAIRMQ socket path to " + socketpath) + os.environ['FAIRMQ_IPC_PREFIX'] = socketpath + + # some maintenance / init work + if args.list_tasks: + print ('List of tasks in this workflow:') + for i,t in enumerate(self.workflowspec['stages'],0): + print (t['name'] + ' (' + str(t['labels']) + ')' + ' ToDo: ' + str(not self.ok_to_skip(i))) + exit (0) + + if args.produce_script != None: + self.produce_script(args.produce_script) + exit (0) + + # execute the user-given global init cmd for this workflow + globalinitcmd = self.globalinit.get("cmd", None) + if globalinitcmd != None: + if not self.execute_globalinit_cmd(globalinitcmd): + exit (1) + + if args.rerun_from: + reruntaskfound=False + for task in self.workflowspec['stages']: + taskname=task['name'] + if re.match(args.rerun_from, taskname): + reruntaskfound=True + taskid=self.tasktoid[taskname] + self.remove_done_flag(find_all_dependent_tasks(self.possiblenexttask, taskid)) + if not reruntaskfound: + print('No task matching ' + args.rerun_from + ' found; cowardly refusing to do anything ') + exit (1) + + # ***************** + # main control loop + # ***************** + candidates = [ tid for tid in self.possiblenexttask[-1] ] + + self.process_list=[] # list of tuples of nodes ids and Popen subprocess instances + + finishedtasks=[] # global list of finished tasks + + try: + + while True: + # sort candidate list according to task weights + candidates = [ (tid, self.taskweights[tid]) for tid in candidates ] + candidates.sort(key=lambda tup: (tup[1][0],-tup[1][1])) # prefer small and same timeframes first then prefer important tasks within frameframe + # remove weights + candidates = [ tid for tid,_ in candidates ] + + finished = [] # --> to account for finished because already done or skipped + actionlogger.debug('Sorted current candidates: ' + str([(c,self.idtotask[c]) for c in candidates])) + self.try_job_from_candidates(candidates, finished) + if len(candidates) > 0 and len(self.process_list) == 0: + self.noprogress_errormsg() + send_webhook(self.args.webhook,"Unable to make further progress: Quitting") + errorencountered = True + break + + finished_from_started = [] # to account for finished when actually started + failing = [] + while self.waitforany(self.process_list, finished_from_started, failing): + if not args.dry_run: + self.monitor(self.process_list) # ---> make async to normal operation? + time.sleep(1) # <--- make this incremental (small wait at beginning) + else: + time.sleep(0.001) + + finished = finished + finished_from_started + actionlogger.debug("finished now :" + str(finished_from_started)) + finishedtasks = finishedtasks + finished + + # perform file cleanup + if self.do_early_file_removal: + self.perform_early_file_removal(finished_from_started) + + if self.is_productionmode: + # we can do some generic cleanup of finished tasks in non-interactive/GRID mode + # TODO: this can run asynchronously + for _t in finished_from_started: + self.production_endoftask_hook(_t) + + # if a task was marked "failed" and we come here (because + # we use --keep-going) ... we need to take out the pid from finished + if len(failing) > 0: + # remove these from those marked finished in order + # not to continue with their children + errorencountered = True + for t in failing: + finished = [ x for x in finished if x != t ] + finishedtasks = [ x for x in finishedtasks if x != t ] + + # if a task was marked as "retry" we simply put it back into the candidate list + if len(self.tids_marked_toretry) > 0: + # we need to remove these first of all from those marked finished + for t in self.tids_marked_toretry: + finished = [ x for x in finished if x != t ] + finishedtasks = [ x for x in finishedtasks if x != t ] + + candidates = candidates + self.tids_marked_toretry + self.tids_marked_toretry = [] + + + # new candidates + for tid in finished: + if self.possiblenexttask.get(tid)!=None: + potential_candidates=list(self.possiblenexttask[tid]) + for candid in potential_candidates: + # try to see if this is really a candidate: + if self.is_good_candidate(candid, finishedtasks) and candidates.count(candid)==0: + candidates.append(candid) + + actionlogger.debug("New candidates " + str( candidates)) + send_webhook(self.args.webhook, "New candidates " + str(candidates)) + + if len(candidates)==0 and len(self.process_list)==0: + break + except Exception as e: + exc_type, exc_obj, exc_tb = sys.exc_info() + fname = os.path.split(exc_tb.tb_frame.f_code.co_filename)[1] + print(exc_type, fname, exc_tb.tb_lineno) + traceback.print_exc() + print ('Cleaning up ') + + self.SIGHandler(0,0) + + endtime = time.perf_counter() + statusmsg = "success" + if errorencountered: + statusmsg = "with failures" + + print ('\n**** Pipeline done ' + statusmsg + ' (global_runtime : {:.3f}s) *****\n'.format(endtime-self.start_time)) + actionlogger.debug("global_runtime : {:.3f}s".format(endtime-self.start_time)) + return errorencountered + + +if args.cgroup!=None: + myPID=os.getpid() + # cgroups such as /sys/fs/cgroup/cpuset//tasks + # or /sys/fs/cgroup/cpu//tasks + command="echo " + str(myPID) + f" > {args.cgroup}" + actionlogger.info(f"Try running in cgroup {args.cgroup}") + waitstatus = os.system(command) + if code := os.waitstatus_to_exitcode(waitstatus): + actionlogger.error(f"Could not apply cgroup") + exit(code) + actionlogger.info("Running in cgroup") + + +# This starts the fanotify fileaccess monitoring process +# if asked for +o2dpg_filegraph_exec = os.getenv("O2DPG_PRODUCE_FILEGRAPH") # switches filegraph monitoring on and contains the executable name +if o2dpg_filegraph_exec: + env = os.environ.copy() + env["FILEACCESS_MON_ROOTPATH"] = os.getcwd() + env["MAXMOTHERPID"] = f"{os.getpid()}" + + fileaccess_log_file_name = f"pipeline_fileaccess_{os.getpid()}.log" + fileaccess_log_file = open(fileaccess_log_file_name, "w") + fileaccess_monitor_proc = subprocess.Popen( + [o2dpg_filegraph_exec], + stdout=fileaccess_log_file, + stderr=subprocess.STDOUT, + env=env) +else: + fileaccess_monitor_proc = None + +try: + # This is core workflow runner invocation + executor=WorkflowExecutor(args.workflowfile,jmax=int(args.maxjobs),args=args) + rc = executor.execute() +finally: + if fileaccess_monitor_proc: + fileaccess_monitor_proc.terminate() # sends SIGTERM + try: + fileaccess_monitor_proc.wait(timeout=5) + except subprocess.TimeoutExpired: + fileaccess_monitor_proc.kill() # force kill if not stopping + # now produce the final filegraph output + o2dpg_root = os.getenv("O2DPG_ROOT") + analyse_cmd = [ + sys.executable, # runs with same Python interpreter + f"{o2dpg_root}/UTILS/FileIOGraph/analyse_FileIO.py", + "--actionFile", actionlogger_file, + "--monitorFile", fileaccess_log_file_name, + "-o", f"pipeline_fileaccess_report_{os.getpid()}.json", + "--basedir", os.getcwd() ] + print (f"Producing FileIOGraph with command {analyse_cmd}") + subprocess.run(analyse_cmd, check=True) + +sys.exit(rc) \ No newline at end of file diff --git a/MC/workflow_runner/o2dpg_runner/README.md b/MC/workflow_runner/o2dpg_runner/README.md new file mode 100644 index 000000000..768ba592b --- /dev/null +++ b/MC/workflow_runner/o2dpg_runner/README.md @@ -0,0 +1,304 @@ +# o2dpg_runner — modular rewrite of the O2DPG workflow runner + +This package is a modular rewrite of the single-file runner +`MC/bin/o2_dpg_workflow_runner.py` (~2000 lines, module-global state). +Both are installed; which one runs is chosen at run time, so the two can be +compared on the same job. Every flag the original parser accepts, this one +accepts, with the same semantics. + +## Where it lives, and how it is selected + +The package and its entry point live in `MC/workflow_runner/`: + +``` +MC/workflow_runner/ + o2dpg_workflow_runner.py # entry point + o2dpg_schedule_simulator.py + o2dpg_runner/ + cli.py # argparse -> RunnerConfig -> Executor + config.py # RunnerConfig dataclass + workflow.py # load / filter / build DAG + graph.py # Kahn topo sort, memoized descendants + resources.py # TaskResources, ResourceManager + monitoring.py # threaded psutil monitor + scheduler/ # timeframe (default), critical_path, best_fit + executor.py # main control loop + cleanup.py # early file removal, log archival + alienv.py # alienv env resolution + cache.py # _done + _done.json fingerprint cache + tests/ +``` + +Nothing calls it directly. `MC/bin/o2_dpg_workflow_runner.py` is a +dispatcher that execs either this runner or the original one, and every +existing call site goes through it unchanged: + +```bash +ALIEN_O2DPG_WORKFLOW_RUNNER=new $O2DPG_ROOT/MC/bin/o2_dpg_workflow_runner.py -f workflow.json +``` + +`legacy` is the default and runs `MC/bin/o2dpg_workflow_runner_legacy.py`. +An unknown value is an error, not a silent fallback. + +No new required dependencies. `psutil` and optionally `graphviz` as before. + +## What changed behaviorally + +The default invocation should reproduce prototype behavior bit-for-bit +(same scheduling decisions, same output files, same log formats). +Everything new is opt-in. + +### New CLI flags (all default to current behavior) + +| Flag | Default | Effect | +| ----------------------------- | ------------- | ------------------------------------------------------------------------------ | +| `--scheduler-policy` | `timeframe` | Choose: `timeframe` (legacy), `critical-path`, or `best-fit`. | +| `--drop-should-break` | off | In `timeframe`, let light tasks slip past a non-fitting heavy task. | +| `--monitor-interval-cpu` | `1.0` (s) | CPU polling cadence for the background monitor thread. | +| `--monitor-interval-mem` | `5.0` (s) | PSS polling cadence (much cheaper to read less often). | +| `--monitor-backend` | `psutil` | Reserved for a future cgroup-v2 backend. | +| `--cache-policy` | `off` | Task-completion cache: `off` (legacy), `lenient`, `strict`. See below. | + +### Removed flags + +- `--webhook` — the debug Mattermost channel integration is gone. The + flag is still accepted (for compatibility) but ignored. +- `--checkpoint-on-failure` — the tarball + `alien.py cp` failure + checkpoint was only used for Grid debugging. Same: accepted, ignored. +- `--cgroup` — superseded by `--systemd-run`. Same: accepted, ignored. + All three log a warning when passed, and none of them aborts, so a JDL + passing one through `ALIEN_O2DPG_ADDITIONAL_WORKFLOW_RUNNER_ARGS` still + runs. + +### Monitoring + +The old runner polled psutil synchronously in the main scheduling loop, +costing 10–20 % of one core on realistic workflows. The new monitor runs +in a background thread with two independent cadences — CPU (cheap, +~1 Hz) and PSS (expensive, ~0.2 Hz). The scheduler reads the latest +snapshot non-blockingly. Result: the runner's self-CPU drops to ~1-2 %. + +The scheduler's dynamic-resource sampling (`--dynamic-resources`) still +triggers in `ResourceManager.unbook()`, preserving ordering with +respect to task completions. + +### Scheduler policies + +Three policies ship, switchable via `--scheduler-policy`: + +- **`timeframe`** (default) — exact prototype behavior. Sort by + `(timeframe, -num_descendants)`; first non-fitting task in the default + pass breaks the pass (this is the legacy quirk that blocks light tasks + behind heavy ones). Set `--drop-should-break` to disable that quirk. +- **`critical-path`** — sort by longest-remaining CPU-weighted path to a + leaf. Standard HEFT-style heuristic; tends to win when resource + estimates are accurate (after `--update-resources`). +- **`best-fit`** — iterative best-fit bin-packing over the candidate set + against the remaining CPU/MEM budget. Maximizes parallelism at the cost + of slightly less predictable task ordering. + +Comparing these three on the same workflow with the same +estimates gives a direct A/B measurement of scheduling strategy impact. + +### Cache policy + +`_done` files remain the primary skip marker (O2 taskwrapper compatibility +is preserved). With `--cache-policy lenient`, a sidecar `_done.json` is +written containing a fingerprint of (command, env subset, software tag, +needs). On rerun, if the command or `needs` list changed, the `_done` +file is removed and the task re-runs. `strict` additionally invalidates +on env/software changes. + +No behavior change unless you pass the flag. + +## What to know + +Things that only came out of running this, and that cost time to find again. + +**A task's first CPU reading is always 0.** `psutil.cpu_percent(interval=None)` +has no baseline on its first call and returns `0.0` by construction. That is +why `sample_resources()` refuses to learn from fewer than three samples, and +why a sampled CPU of `0` is treated as missing information rather than as a +task that needs no CPU. Propagating a zero would make every sibling look free +and admit them all at once. + +**A monitor tick is not a poll.** The wait loop polls at 0.1 s ramping to 1 s +while the monitor thread fires at `--monitor-interval-cpu`. Samples must be +taken once per tick; recording per poll grows the sample lists without bound +over a long job and, worse, satisfies the three-sample guard with copies of +one reading. + +**A tracer or a monitor has to sit inside the systemd scope.** With +`--systemd-run`, wrapping the outside of `systemd-run` means observing +`systemd-run` and nothing else. The wrapper goes on the inner command. + +**`--maxjobs 1` serialises a workflow; `--cpu-limit 1` does not.** The latter +makes every task that declares more than one core unschedulable, and the run +stops with `ResourceLimitExceeded`. + +**The prototype and this runner disagree on short tasks, in this runner's +favour.** The prototype samples from its scheduling loop at roughly 5 s, so a +task shorter than ~15 s never reaches three samples and `--dynamic-resources` +learns nothing from it. The monitor thread here runs at 1 s. + +## Bug fixes incorporated + +Silently fixed relative to the prototype: + +1. `TaskResources.is_within_limits()` compared CPU to mem_limit instead + of MEM to mem_limit; the memory safety net was effectively disabled. +2. `find_all_dependent_tasks()` cached duplicates but returned + deduped; cache hits returned different values than cache misses. +3. `filter_workflow()` aliased the caller's dict and mutated it in place. +4. `getallrequirements()` recursed without memoization — exponential on + diamond DAGs; `sys.setrecursionlimit(100000)` was a workaround. +5. `send_webhook()` shell-interpolated task names into `os.system` + (command injection). Removed along with the webhook feature. +6. `SIGHandler` only caught SIGINT; Grid preemption via SIGTERM was + ignored. Now both are handled. +7. The emitted `produce_script` output used `cd $OLDPWD` which broke if + a task `cd`s internally. Now uses subshells: `( cd "$workdir" && ... )`. +8. `candidates` list used `.count()` for membership checks (O(n)); + replaced by set-based lookups. +9. A sampled CPU of `0` was propagated to the un-started sibling tasks, + which then all looked free and were admitted at once. It is now treated + as missing information, as a sampled memory of `0` already was. The + prototype carries this too, masked by its slower monitoring. +10. The dry-run path returned a plain `subprocess.Popen`, which has no + `.nice()`. Fixed here the same way it was later fixed in the + prototype: by returning a `psutil.Popen`. + +## Running the tests + +```bash +cd MC/workflow_runner +python -m pytest o2dpg_runner/tests/ -q +``` + +`pytest` is required. The tests run in CI, in the `Workflow-runner unit +tests` job of `.github/workflows/syntax-checks.yml`. + +The tests cover: +- `test_graph.py` — Kahn topological sort, memoized descendants/ancestors, + longest path, diamond + deep-chain cases. +- `test_workflow.py` — load, global-init extraction, filtering by target + and by label, regex, resource-estimate update. +- `test_resources.py` — booking/unbooking, semaphores, related-task + grouping, dynamic sampling, limit enforcement. +- `test_scheduler.py` — all three policies, the `should_break` quirk + and its removal, `n_backfill_max` cap, semaphore blocking. +- `test_cache.py` — cache policies (off/lenient/strict), fingerprint + sensitivity, sidecar round-trip. +- `test_executor_e2e.py` — the tiny fixture workflow driven end-to-end + with real subprocesses, exercising each policy, `--dry-run`, + `--produce-script`, rerun-from-cache behavior. +- `test_simulator.py` — simulator-only coverage for Amdahl-derived + critical-path weights, unschedulable-task handling, and simulated + backfill behaviour (`slowdown` and `holefill`). + +Integration test (from the prototype, still valid): +```bash +NSIGEVENTS=5 NTIMEFRAMES=2 bash MC/bin/tests/wf_test_pp.sh +``` + +## A/B measurement + +The Python entry point accepts the same `workflow.json` under all +scheduling policies. A typical comparison: + +```bash +# baseline (prototype-equivalent) +./o2dpg_workflow_runner.py -f wf.json \ + --metric-logfile metric_timeframe.log + +# drop the should_break quirk +./o2dpg_workflow_runner.py -f wf.json --drop-should-break \ + --metric-logfile metric_timeframe_nobrk.log + +# critical path +./o2dpg_workflow_runner.py -f wf.json --scheduler-policy critical-path \ + --metric-logfile metric_cp.log + +# best-fit bin-packing +./o2dpg_workflow_runner.py -f wf.json --scheduler-policy best-fit \ + --metric-logfile metric_bf.log +``` + +The metric logs have the same schema as before (`o2dpg_sim_metrics.py` +post-processing is unaffected), plus the run's meta line now records +`scheduler_policy`, `drop_should_break`, and `cache_policy` for easy +downstream grouping. + +## Simulator notes + +`MC/bin/o2dpg_schedule_simulator.py` is an offline discrete-event model +of the runner. It exists to compare policies and tune worker-count / +resource parameters quickly, not to emulate Linux scheduling perfectly. + +- The simulator now uses the same walltime-weighted critical-path input + as the runner when learned `resources.walltime` data is available. +- Amdahl worker overrides are applied before simulator scheduler-state + construction, so optimization runs evaluate policies against the same + task costs they simulate. +- Tasks that exceed the hard simulated CPU/MEM limits are kept in the + workflow model and reported as unschedulable, rather than being + dropped from resource bookkeeping. + +### Simulated backfill + +The real runner has a two-lane admission model: default tasks stay +within the hard budget, while backfill tasks may use a bounded amount of +overcommit and run at lower priority. The simulator now offers a +"sweet-spot" family of approximations for that behaviour: + +- `--backfill-model off` — no backfill, one hard budget only. +- `--backfill-model structural` — replay the runner's second admission + lane (`n_backfill`, CPU factor, MEM factor), but do not change task + duration. +- `--backfill-model slowdown` — same structural backfill admission, plus + a single fitted slowdown factor applied to backfill task walltimes. +- `--backfill-model holefill` — preferred realistic mode. Foreground + tasks keep their nominal duration; backfill tasks consume only the CPU + left idle by currently running foreground tasks. When the foreground + hole is smaller than the task's nominal CPU demand, the task slows + down proportionally. When the hole is large enough, it runs at nominal + speed. + +Relevant knobs: + +- `--n-backfill` +- `--backfill-cpu-factor` +- `--backfill-mem-factor` +- `--backfill-slowdown-factor` + +`holefill` is the recommended mode for scheduling studies: it +keeps simulated CPU efficiency physically bounded by 100%, tracks +observed runtime improvements from backfilling much better than the +single-factor slowdown model, and still stays simple enough to explain. + +Two modelling choices are worth keeping in mind: + +- foreground tasks are assumed not to slow down due to backfill; +- hole allocation is online and greedy, so enabling backfill may still + change the order in which later tasks become runnable. + +This is intentionally a scheduler-level approximation, not a kernel CPU +sharing model. It is accurate enough for comparative studies while still +remaining easy to reason about and calibrate against real runs. + +### Suggested simulator usage + +For realistic policy comparison, use learned resources and the holefill +backfill model: + +```bash +./o2dpg_schedule_simulator.py \ + --timeframes 1 2 4 5 8 12 20 \ + --update-resources learned.json \ + -f workflow.json \ + --backfill-model holefill +``` + +Use `off` as the baseline and `holefill` as the realistic backfill +comparison. The older `slowdown` mode remains useful as a coarse control +study, but it is no longer the preferred setting for reporting numbers. diff --git a/MC/workflow_runner/o2dpg_runner/__init__.py b/MC/workflow_runner/o2dpg_runner/__init__.py new file mode 100644 index 000000000..0c8a58525 --- /dev/null +++ b/MC/workflow_runner/o2dpg_runner/__init__.py @@ -0,0 +1,11 @@ +""" +o2dpg_runner: modular rewrite of the O2DPG workflow runner. + +Originally started February 2021 by sandro.wenzel@cern.ch as a single-file +prototype (MC/bin/o2dpg_workflow_runner.py). This package is the modular +refactor, keeping full CLI and behavioral compatibility by default while +exposing pluggable scheduling policies, a threaded resource monitor, and +a cleaner testable structure. +""" + +__version__ = "2.0.0" diff --git a/MC/workflow_runner/o2dpg_runner/alienv.py b/MC/workflow_runner/o2dpg_runner/alienv.py new file mode 100644 index 000000000..692e0f077 --- /dev/null +++ b/MC/workflow_runner/o2dpg_runner/alienv.py @@ -0,0 +1,66 @@ +"""Resolving alternative software environments via alienv.""" + +from __future__ import annotations + +import logging +import os +import subprocess +from typing import Dict, Optional + +log = logging.getLogger(__name__) + + +def _load_env_file(path: str) -> Dict[str, str]: + env: Dict[str, str] = {} + with open(path, "r") as f: + for raw in f: + line = raw.strip() + if not line or line.startswith("#"): + continue + if line.startswith("declare -x "): + line = line.replace("declare -x ", "", 1) + if "=" not in line: + env[line.strip()] = "" + else: + k, v = line.split("=", 1) + env[k.strip()] = v.strip('"') + return env + + +def get_alienv_software_environment(packagestring: Optional[str]) -> Dict[str, str]: + """Resolve ``packagestring`` to an env dict. + + Accepts: + - None / '' / 'None' -> empty dict + - a path to a file -> 'export > env.txt' format, parsed + - an alienv spec -> calls /cvmfs/alice.cern.ch/bin/alienv printenv + """ + if not packagestring or packagestring == "None": + return {} + + if os.path.exists(packagestring) and os.path.isfile(packagestring): + log.info("Taking software environment from file %s", packagestring) + return _load_env_file(packagestring) + + cmd = "/cvmfs/alice.cern.ch/bin/alienv printenv " + packagestring + proc = subprocess.Popen( + [cmd], stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True + ) + envstring, err = proc.communicate() + if err: + err_s = err.decode() + if err_s.strip(): + print(err_s) + raise RuntimeError(f"alienv printenv failed for {packagestring}") + + envmap: Dict[str, str] = {} + for t in envstring.decode().split(";"): + if "=" in t: + k, v = t.rstrip().split("=", 1) + envmap[k] = v + elif "export" in t: + tokens = t.split() + if len(tokens) >= 2: + variable = tokens[1] + envmap.setdefault(variable, "") + return envmap diff --git a/MC/workflow_runner/o2dpg_runner/cache.py b/MC/workflow_runner/o2dpg_runner/cache.py new file mode 100644 index 000000000..b7e937f0b --- /dev/null +++ b/MC/workflow_runner/o2dpg_runner/cache.py @@ -0,0 +1,192 @@ +"""Task-completion cache. + +Preserves the O2-taskwrapper-compatible `.log_done` marker file as +the primary source of truth (so that downstream tools and the wrapper +itself continue to work unchanged), and optionally writes a sidecar +`.log_done.json` containing a fingerprint of the inputs that +determine the task's output. + +Cache policies: + off - current behavior: a _done file means "skip" + lenient - read _done.json when present; invalidate only if the + command string changed. Warn on env / software changes. + strict - invalidate on any fingerprint change. + +Fingerprint components: + cmd_hash hash of the task command string + env_hash hash of the allow-listed subset of the task env + software the alienv package string (or '' if default) + needs list of upstream task names + +The sidecar is written best-effort after the _done file exists (so +torn writes don't leave stale fingerprints around). +""" + +from __future__ import annotations + +import hashlib +import json +import logging +import os +from typing import Dict, List, Optional, Tuple + +log = logging.getLogger(__name__) + +# Env vars that actually affect task semantics. Keep this conservative: +# everything else is considered noise (TMPDIR, PWD, terminal colors, etc.). +# Users can override via the workflow by not relying on env and by expressing +# variation through the cmd itself. +SEMANTIC_ENV_KEYS = ( + "ALICE_O2_VERSION", "O2_ROOT", "O2DPG_ROOT", "O2PHYSICS_ROOT", + "NEVENTS", "NWORKERS", "SEED", "INTERACTIONRATE", + "SIMENGINE", "NSIGEVENTS", "NTIMEFRAMES", "GENERATOR", + "ALIDIST_TAG", +) + + +def _hash(*parts: str) -> str: + h = hashlib.sha256() + for p in parts: + h.update(p.encode("utf-8", "replace")) + h.update(b"\x1f") # separator + return h.hexdigest()[:16] + + +def compute_fingerprint( + task: Dict, + alienv_package: str = "", + allow_env_keys: Tuple[str, ...] = SEMANTIC_ENV_KEYS, +) -> Dict[str, str]: + """Compute a fingerprint for a task spec. + + Does NOT depend on upstream task fingerprints -- upstream changes + propagate via the _done/_done.json of upstream tasks (if upstream + was re-run, its _done file was removed and the current task will + see a missing dependency marker and re-run). + """ + cmd = task.get("cmd", "") or "" + env_subset = {} + task_env = task.get("env") or {} + for k in allow_env_keys: + if k in task_env: + env_subset[k] = str(task_env[k]) + env_str = json.dumps(env_subset, sort_keys=True) + needs = sorted(task.get("needs", []) or []) + needs_str = json.dumps(needs) + + return { + "cmd_hash": _hash(cmd), + "env_hash": _hash(env_str), + "software": alienv_package or "", + "needs": needs, + "cmd_preview": cmd[:120], # for human debugging + } + + +def done_path(logfile: str) -> str: + return logfile + "_done" + + +def fingerprint_path(logfile: str) -> str: + return logfile + "_done.json" + + +class TaskCache: + """Policy-aware interface for checking/recording task completion.""" + + def __init__(self, policy: str = "off"): + if policy not in ("off", "lenient", "strict"): + raise ValueError(f"unknown cache policy: {policy}") + self.policy = policy + + def is_done(self, logfile: str, current_fp: Dict[str, str]) -> bool: + """Return True iff the task can be skipped. + + logfile: path like /cwd/taskname.log (we append _done / _done.json) + current_fp: fingerprint dict from compute_fingerprint(task) + """ + dp = done_path(logfile) + if not (os.path.exists(dp) and os.path.isfile(dp)): + return False + + if self.policy == "off": + return True + + fp_path = fingerprint_path(logfile) + if not os.path.exists(fp_path): + # Old run; nothing to compare against. + if self.policy == "strict": + log.info("%s: strict cache policy but no fingerprint -> invalidating", logfile) + self._invalidate(logfile) + return False + log.debug("%s: no fingerprint sidecar; keeping _done (lenient)", logfile) + return True + + try: + with open(fp_path) as f: + prev = json.load(f) + except (OSError, json.JSONDecodeError) as e: + log.warning("%s: fingerprint unreadable (%s); invalidating", fp_path, e) + self._invalidate(logfile) + return False + + cmd_changed = prev.get("cmd_hash") != current_fp["cmd_hash"] + env_changed = prev.get("env_hash") != current_fp["env_hash"] + sw_changed = prev.get("software") != current_fp["software"] + needs_changed = prev.get("needs") != current_fp["needs"] + + if self.policy == "lenient": + if cmd_changed or needs_changed: + log.info("%s: cmd/needs changed -> invalidating (lenient)", logfile) + self._invalidate(logfile) + return False + if env_changed: + log.warning("%s: env fingerprint changed but keeping cache (lenient)", logfile) + if sw_changed: + log.warning("%s: software fingerprint changed but keeping cache (lenient)", + logfile) + return True + + # strict + if cmd_changed or env_changed or sw_changed or needs_changed: + reasons = [] + if cmd_changed: reasons.append("cmd") + if env_changed: reasons.append("env") + if sw_changed: reasons.append("software") + if needs_changed: reasons.append("needs") + log.info("%s: changed (%s) -> invalidating (strict)", logfile, ",".join(reasons)) + self._invalidate(logfile) + return False + return True + + def record(self, logfile: str, current_fp: Dict[str, str]) -> None: + """Write the fingerprint sidecar after the task's _done file exists. + + Best effort: errors are logged but not raised. + """ + if self.policy == "off": + return + dp = done_path(logfile) + if not os.path.exists(dp): + # _done doesn't exist (e.g. skipped in dry-run); nothing to record. + return + try: + with open(fingerprint_path(logfile), "w") as f: + json.dump(current_fp, f, indent=2) + except OSError as e: + log.warning("Could not write fingerprint for %s: %s", logfile, e) + + def _invalidate(self, logfile: str) -> None: + for p in (done_path(logfile), fingerprint_path(logfile)): + try: + if os.path.exists(p): + os.remove(p) + except OSError as e: + log.warning("Could not remove %s: %s", p, e) + + +def remove_done_flag(logfile: str) -> None: + """Explicit invalidation used by --rerun-from.""" + for p in (done_path(logfile), fingerprint_path(logfile)): + if os.path.exists(p) and os.path.isfile(p): + os.remove(p) diff --git a/MC/workflow_runner/o2dpg_runner/cleanup.py b/MC/workflow_runner/o2dpg_runner/cleanup.py new file mode 100644 index 000000000..a5e526ea5 --- /dev/null +++ b/MC/workflow_runner/o2dpg_runner/cleanup.py @@ -0,0 +1,221 @@ +"""Cleanup utilities invoked after a task completes. + +Two concerns: + 1. Early file removal based on a FileIOGraph-derived file dependency map + (--remove-files-early). Files that no task will read/write again are + deleted to keep disc pressure low during large productions. + + 2. Production-mode log archival: .log / .log_done / .log_time files are + appended to a tar archive and removed. Mirrors the prototype's + production_endoftask_hook(). + +Both are side-effecting; errors are logged but don't abort the run. +""" + +from __future__ import annotations + +import json +import logging +import os +import re +import tarfile +from typing import Dict, List, Optional, Set + +log = logging.getLogger(__name__) + +_TF_PATH_RE = re.compile(r"^\./tf(?P\d+)/") + + +def _task_template_for_timeframe(task_name: str, source_tf: int) -> str: + """Turn task_7 from a tf7 observation into task_{tf}. + + Global tasks such as aodmerge, and any explicit cross-timeframe task + references, are left untouched. + """ + suffix = f"_{source_tf}" + if task_name.endswith(suffix): + return f"{task_name[:-len(suffix)]}_{{tf}}" + return task_name + + +def _task_template_from_placeholder(task_name: str) -> str: + return re.sub(r"_X$", "_{tf}", task_name) + + +def _filegraph_expand_timeframes( + data: Dict, + timeframes: Set[int], + target_namelist: List[str], + logger: Optional[logging.Logger] = None, +) -> Dict[str, List[Dict]]: + """Build canonical per-timeframe file templates and replicate per TF. + + The FileIOGraph may have been recorded with one or many timeframes. We + merge all observed timeframe-local entries by normalising ``./tfX/...`` to + ``./tf{tf}/...`` and task suffixes ``_X`` to ``_{tf}``, then instantiate the + merged template for the timeframes of the current workflow. + """ + logger = logger if logger is not None else log + templates: Dict[str, Dict] = {} + source_tfs: Set[int] = set() + + template_report = data.get("file_template_report") or [] + if template_report: + for entry in template_report: + filename = entry.get("file", "") + if filename.startswith("./tfX/"): + file_template = filename.replace("./tfX/", "./tf{tf}/", 1) + elif "./tf{tf}/" in filename: + file_template = filename + else: + continue + + source_tfs.update(int(tf) for tf in entry.get("source_timeframes", [])) + merged = templates.setdefault( + file_template, + { + "file": file_template, + "written_by": set(), + "read_by": set(), + "keep": bool(entry.get("keep", False)), + }, + ) + merged["keep"] = merged["keep"] or bool(entry.get("keep", False)) + for task in entry.get("written_by", []): + merged["written_by"].add(_task_template_from_placeholder(task)) + for task in entry.get("read_by", []): + merged["read_by"].add(_task_template_from_placeholder(task)) + else: + for entry in data.get("file_report", []): + filename = entry.get("file", "") + match = _TF_PATH_RE.match(filename) + if match is None: + continue + + source_tf = int(match.group("tf")) + source_tfs.add(source_tf) + file_template = _TF_PATH_RE.sub("./tf{tf}/", filename, count=1) + merged = templates.setdefault( + file_template, + { + "file": file_template, + "written_by": set(), + "read_by": set(), + "keep": bool(entry.get("keep", False)), + }, + ) + merged["keep"] = merged["keep"] or bool(entry.get("keep", False)) + for task in entry.get("written_by", []): + merged["written_by"].add(_task_template_for_timeframe(task, source_tf)) + for task in entry.get("read_by", []): + merged["read_by"].add(_task_template_for_timeframe(task, source_tf)) + + if not templates: + logger.warning("FileIOGraph contains no ./tfN/ file entries; early removal disabled") + return {} + + logger.info( + "FileIOGraph timeframe template built from observed timeframe(s) %s: %d file pattern(s)", + sorted(source_tfs), + len(templates), + ) + + result: Dict[str, List[Dict]] = {} + for i in timeframes: + if i == -1: + continue + new_entries: List[Dict] = [] + for template in templates.values(): + written_by = sorted(t.format(tf=i) for t in template["written_by"]) + read_by = sorted(t.format(tf=i) for t in template["read_by"]) + expanded = { + "file": template["file"].format(tf=i), + "written_by": written_by, + "read_by": read_by, + } + if template["keep"] or any(w in target_namelist for w in written_by): + expanded["keep"] = True + new_entries.append(expanded) + result[f"timeframe-{i}"] = new_entries + return result + + +class EarlyFileRemover: + """Owns the timeframe-expanded file dependency dict and performs + per-task-completion file deletion.""" + + def __init__( + self, + filegraph_path: str, + timeframes: Set[int], + target_namelist: List[str], + logger: Optional[logging.Logger] = None, + ): + self.log = logger if logger is not None else log + with open(filegraph_path) as f: + data = json.load(f) + self.file_dict = _filegraph_expand_timeframes( + data, timeframes, target_namelist, logger=self.log, + ) + self.timeframes = timeframes + # Pre-build a reverse index: task_name -> list[file_entry dict] (all TFs) + # so completions don't re-scan the whole file map. + self._by_task: Dict[str, List[Dict]] = {} + for entries in self.file_dict.values(): + for e in entries: + for t in e.get("written_by", []): + self._by_task.setdefault(t, []).append(e) + for t in e.get("read_by", []): + self._by_task.setdefault(t, []).append(e) + + def on_task_done(self, taskname: str) -> None: + entries = self._by_task.get(taskname, []) + for entry in entries: + if taskname in entry.get("read_by", []): + entry["read_by"].remove(taskname) + if taskname in entry.get("written_by", []): + entry["written_by"].remove(taskname) + if (not entry.get("read_by") and not entry.get("written_by") + and not entry.get("keep", False)): + self._remove_if_exists(entry["file"]) + + def _remove_if_exists(self, path: str) -> bool: + if os.path.exists(path): + try: + sz = os.path.getsize(path) + os.remove(path) + self.log.info("Removing %s (no longer needed); freed %.2f MB", + path, sz / 1024.0 / 1024.0) + return True + except OSError as e: + self.log.warning("Could not remove %s: %s", path, e) + return False + + +def archive_task_logs(logfile: str, logger: Optional[logging.Logger] = None) -> None: + """Append , _done, _time to a tar archive + and delete the originals. Used in production mode.""" + logger = logger if logger is not None else log + done = logfile + "_done" + timef = logfile + "_time" + try: + tf = tarfile.open(name="pipeline_log_archive.log.tar", mode="a") + except Exception as e: + logger.warning("Could not open log archive: %s", e) + return + try: + for path in (logfile, done, timef): + if os.path.exists(path): + try: + tf.add(path) + except Exception as e: + logger.warning("tar add %s failed: %s", path, e) + finally: + tf.close() + + for path in (logfile, done, timef): + if os.path.exists(path): + try: + os.remove(path) + except OSError as e: + logger.warning("Could not remove %s: %s", path, e) diff --git a/MC/workflow_runner/o2dpg_runner/cli.py b/MC/workflow_runner/o2dpg_runner/cli.py new file mode 100644 index 000000000..defab9ad2 --- /dev/null +++ b/MC/workflow_runner/o2dpg_runner/cli.py @@ -0,0 +1,468 @@ +"""Command-line entry point. + +Translates argparse -> RunnerConfig, builds loggers, creates the +WorkflowExecutor, and invokes it. All semantics of the original script +are preserved; new flags are additive and default-compatible. +""" + +from __future__ import annotations + +import argparse +import json +import logging +import os +import shutil +import subprocess +import sys +from typing import Optional, Tuple + +import psutil + +from .config import RunnerConfig +from .workflow import build_workflow, load_json +from .executor import WorkflowExecutor + +_FORMATTER = logging.Formatter("%(asctime)s %(levelname)s %(message)s") +_IN_SLICE_ENV = "O2DPG_RUNNER_IN_SLICE" + + +def _setup_logger(name: str, logfile: str, level: int = logging.INFO) -> logging.Logger: + handler = logging.FileHandler(logfile, mode="w") + handler.setFormatter(_FORMATTER) + logger = logging.getLogger(name) + logger.setLevel(level) + logger.handlers.clear() + logger.addHandler(handler) + logger.propagate = False + return logger + + +def build_parser() -> argparse.ArgumentParser: + max_system_mem = psutil.virtual_memory().total + default_mem = 0.9 * max_system_mem / 1024.0 / 1024.0 + + p = argparse.ArgumentParser( + description="Parallel execution of an O2-DPG data/job DAG under resource constraints.", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + p.add_argument("-f", "--workflowfile", required=True) + p.add_argument("-jmax", "--maxjobs", type=int, default=100) + p.add_argument("-k", "--keep-going", action="store_true") + p.add_argument("--dry-run", action="store_true") + p.add_argument("--visualize-workflow", action="store_true") + p.add_argument("--target-labels", nargs="+", default=[]) + p.add_argument("-tt", "--target-tasks", nargs="+", default=["*"]) + p.add_argument("--produce-script", default=None) + p.add_argument("--rerun-from", default=None) + p.add_argument("--list-tasks", action="store_true") + + # Resources + p.add_argument("--update-resources", dest="update_resources", default=None) + p.add_argument("--dynamic-resources", dest="dynamic_resources", action="store_true") + p.add_argument("--optimistic-resources", dest="optimistic_resources", action="store_true") + p.add_argument("--n-backfill", dest="n_backfill", type=int, default=1) + p.add_argument("--mem-limit", type=float, default=default_mem, help="in MB") + p.add_argument("--cpu-limit", type=float, default=8) + + # systemd-run slice confinement (replaces the old --cgroup option) + p.add_argument( + "--systemd-run", + dest="systemd_run_spec", + default=None, + metavar="SPEC", + help=( + "Relaunch the whole runner (and all child processes) inside a transient " + "systemd scope unit with resource limits. SPEC is a slash-separated list " + "of key:value pairs. Supported keys: ncpus (number of CPU cores, e.g. 8), " + "mem (memory limit in systemd format, e.g. 16G or 16384M). " + "Either key may be omitted. Examples: \"ncpus:8/mem:16G\", \"ncpus:4\", \"mem:32G\"." + ), + ) + + # Scheduling (new) + p.add_argument("--scheduler-policy", default="timeframe", + choices=["timeframe", "critical-path", "best-fit"]) + p.add_argument("--drop-should-break", action="store_true", + help="In timeframe policy, don't stop scanning on the first " + "non-fitting task (lets light tasks slip past heavy ones).") + + # Monitoring (new) + p.add_argument("--monitor-interval-cpu", type=float, default=1.0) + p.add_argument("--monitor-interval-mem", type=float, default=1.0) + p.add_argument("--monitor-backend", default="psutil", choices=["psutil"]) + + # Cache (new) + p.add_argument("--cache-policy", default="off", + choices=["off", "lenient", "strict"]) + + # Control + p.add_argument("--stdout-on-failure", action="store_true") + p.add_argument("--retry-on-failure", type=int, default=0) + p.add_argument("--no-rootinit-speedup", action="store_true") + p.add_argument("--remove-files-early", type=str, default="") + + # Accept-and-ignore for backward compatibility of call sites + # that still pass these flags. They have no effect. + p.add_argument("--webhook", default=None, help=argparse.SUPPRESS) + p.add_argument("--checkpoint-on-failure", default=None, help=argparse.SUPPRESS) + # superseded by --systemd-run; a JDL still passing it must not abort here + p.add_argument("--cgroup", default=None, help=argparse.SUPPRESS) + + # Logging + p.add_argument("--action-logfile", default=None) + p.add_argument("--metric-logfile", default=None) + p.add_argument("--production-mode", action="store_true") + + return p + + +def _args_to_config(ns: argparse.Namespace) -> RunnerConfig: + target_tasks = [f.strip('"').strip("'") for f in ns.target_tasks] + # Extract slice name from spec so the executor can name child scopes. + slice_name: Optional[str] = None + if ns.systemd_run_spec: + try: + _, _, slice_name = _parse_systemd_run_spec(ns.systemd_run_spec) + except ValueError: + pass # error already caught at re-exec time + return RunnerConfig( + workflowfile=ns.workflowfile, + maxjobs=ns.maxjobs, + mem_limit=ns.mem_limit, + cpu_limit=ns.cpu_limit, + n_backfill=ns.n_backfill, + update_resources=ns.update_resources, + dynamic_resources=ns.dynamic_resources, + optimistic_resources=ns.optimistic_resources, + in_systemd_slice=bool(os.environ.get(_IN_SLICE_ENV)), + systemd_run_spec=ns.systemd_run_spec, + systemd_slice_name=slice_name, + scheduler_policy=ns.scheduler_policy, + drop_should_break=ns.drop_should_break, + monitor_interval_cpu=ns.monitor_interval_cpu, + monitor_interval_mem=ns.monitor_interval_mem, + monitor_backend=ns.monitor_backend, + cache_policy=ns.cache_policy, + target_tasks=target_tasks, + target_labels=list(ns.target_labels), + keep_going=ns.keep_going, + dry_run=ns.dry_run, + visualize_workflow=ns.visualize_workflow, + produce_script=ns.produce_script, + rerun_from=ns.rerun_from, + list_tasks=ns.list_tasks, + retry_on_failure=ns.retry_on_failure, + no_rootinit_speedup=ns.no_rootinit_speedup, + remove_files_early=ns.remove_files_early, + stdout_on_failure=ns.stdout_on_failure, + production_mode=ns.production_mode, + action_logfile=ns.action_logfile, + metric_logfile=ns.metric_logfile, + ) + + +def _parse_systemd_run_spec(spec: str) -> Tuple[Optional[str], Optional[str], str]: + """Parse "ncpus:N/mem:M/name:S" into (cpu_quota_str, mem_str, slice_name). + + ncpus is given as a number of cores and converted to systemd CPUQuota + format (e.g. 8 cores → "800%"). mem is passed through as-is. + name sets the systemd slice name (default "o2dpg"). Any part may be absent. + """ + cpu_quota: Optional[str] = None + mem: Optional[str] = None + slice_name: str = "o2dpg" + for part in spec.split("/"): + part = part.strip() + if not part: + continue + if ":" not in part: + raise ValueError(f"Expected key:value in --systemd-run spec, got: {part!r}") + key, _, val = part.partition(":") + key = key.strip().lower() + val = val.strip() + if key == "ncpus": + try: + cores = float(val) + except ValueError: + raise ValueError(f"ncpus must be a number, got: {val!r}") + cpu_quota = f"{int(cores * 100)}%" + elif key == "mem": + mem = val + elif key == "name": + slice_name = val + else: + raise ValueError(f"Unknown key in --systemd-run spec: {key!r}. " + f"Supported: ncpus, mem, name") + return cpu_quota, mem, slice_name + + +def _parse_mem_to_bytes(mem_str: str) -> Optional[int]: + """Convert a memory size string (e.g. "16G", "512M") to bytes.""" + suffixes = {"K": 1024, "M": 1024**2, "G": 1024**3, "T": 1024**4} + s = mem_str.strip() + if s and s[-1].upper() in suffixes: + try: + return int(float(s[:-1]) * suffixes[s[-1].upper()]) + except ValueError: + pass + try: + return int(s) + except ValueError: + return None + + +def _apply_slice_cgroup_limits( + cpu_quota: Optional[str], + mem_str: Optional[str], + logger: logging.Logger, +) -> None: + """Write resource limits to the parent slice's cgroup directory. + + Called after re-exec inside the slice. The runner's own scope lives at + .../kaz1.slice/o2dpg-runner-.scope/; one dirname() up is the slice + that covers both the runner and all sibling per-task scopes. + + This is the portable fallback for systemd < 246 which lacks + --slice-property. Writing directly to cpu.max / memory.max works on + any cgroup v2 system where the user owns the cgroup. + """ + cgroup_rel: Optional[str] = None + try: + with open("/proc/self/cgroup") as fh: + for line in fh: + parts = line.strip().split(":", 2) + if len(parts) == 3 and parts[0] == "0": + cgroup_rel = parts[2].lstrip("/") + break + except OSError: + pass + + if not cgroup_rel: + logger.warning("Cannot determine cgroup path; slice limits not applied.") + return + + scope_dir = f"/sys/fs/cgroup/{cgroup_rel}" + slice_dir = os.path.dirname(scope_dir) + + if cpu_quota: + pct = float(cpu_quota.rstrip("%")) + # cgroup cpu.max format: " " + # 100% = 1 core = 100000 usec per 100000 usec period + quota_usec = int(pct * 1000) + cpu_max = os.path.join(slice_dir, "cpu.max") + try: + with open(cpu_max, "w") as fh: + fh.write(f"{quota_usec} 100000\n") + logger.info("Slice CPUQuota=%s applied → %s", cpu_quota, cpu_max) + except OSError as e: + logger.warning("Could not set CPUQuota on slice (%s): %s", cpu_max, e) + + if mem_str: + mem_bytes = _parse_mem_to_bytes(mem_str) + if mem_bytes is not None: + mem_max = os.path.join(slice_dir, "memory.max") + try: + with open(mem_max, "w") as fh: + fh.write(f"{mem_bytes}\n") + logger.info("Slice MemoryMax=%s (%d bytes) applied → %s", + mem_str, mem_bytes, mem_max) + except OSError as e: + logger.warning("Could not set MemoryMax on slice (%s): %s", mem_max, e) + + +def _maybe_reexec_in_slice(ns: argparse.Namespace) -> None: + """If --systemd-run is set and we are not already in the slice, re-exec. + + Uses os.execvp so the current process image is replaced by systemd-run, + which creates a transient scope cgroup and then exec's the runner again. + The child runner sees O2DPG_RUNNER_IN_SLICE=1 and skips this function. + """ + spec = getattr(ns, "systemd_run_spec", None) + if not spec: + return + if os.environ.get(_IN_SLICE_ENV): + return # already inside the slice + + if not shutil.which("systemd-run"): + print( + "Warning: --systemd-run requested but systemd-run not found on PATH; " + "continuing without slice confinement.", + file=sys.stderr, + ) + return + + try: + _, _, slice_name = _parse_systemd_run_spec(spec) + except ValueError as e: + print(f"Error in --systemd-run spec: {e}", file=sys.stderr) + sys.exit(1) + + # Ensure the slice name has the .slice suffix expected by systemd. + systemd_slice = slice_name if slice_name.endswith(".slice") else f"{slice_name}.slice" + unit_name = f"o2dpg-runner-{os.getpid()}.scope" + cmd = ["systemd-run", "--user", "--scope", "--collect", + f"--unit={unit_name}", f"--slice={systemd_slice}"] + # Resource limits are NOT passed here; they are written directly to the + # slice cgroup after re-exec via _apply_slice_cgroup_limits(). + # --property=CPUQuota applies only to the scope (runner), not to the + # sibling task scopes. --slice-property would be correct but requires + # systemd ≥ 246. Direct cgroup writes work on all versions. + cmd += ["--", sys.executable] + sys.argv + + os.environ[_IN_SLICE_ENV] = "1" + try: + os.execvp(cmd[0], cmd) + except OSError as e: + # execvp only returns on failure + del os.environ[_IN_SLICE_ENV] + print( + f"Warning: could not exec systemd-run ({e}); " + "continuing without slice confinement.", + file=sys.stderr, + ) + + +def _maybe_draw_workflow(raw_spec): + try: + from graphviz import Digraph + except ImportError: + print("graphviz not installed; cannot draw workflow") + return + dot = Digraph(comment="MC workflow") + name_to_idx = {} + for i, node in enumerate(raw_spec["stages"]): + name_to_idx[node["name"]] = i + dot.node(str(i), node["name"]) + for node in raw_spec["stages"]: + to_i = name_to_idx[node["name"]] + for r in node.get("needs", []): + if r in name_to_idx: + dot.edge(str(name_to_idx[r]), str(to_i)) + dot.render("workflow.gv") + + +def _launch_fileaccess_sidecar(actionlogger_file: str): + """Start the fanotify-based file-IO graph sidecar if requested.""" + exe = os.getenv("O2DPG_PRODUCE_FILEGRAPH") + if not exe: + return None, None, None + env = os.environ.copy() + env["FILEACCESS_MON_ROOTPATH"] = os.getcwd() + env["MAXMOTHERPID"] = f"{os.getpid()}" + log_file = f"pipeline_fileaccess_{os.getpid()}.log" + fh = open(log_file, "w") + proc = subprocess.Popen( + [exe], stdout=fh, stderr=subprocess.STDOUT, env=env, + ) + return proc, fh, log_file + + +def main(argv=None) -> int: + ns = build_parser().parse_args(argv) + _maybe_reexec_in_slice(ns) # may replace this process; returns only if not re-execing + cfg = _args_to_config(ns) + + # loggers + action_log = cfg.action_logfile or f"pipeline_action_{os.getpid()}.log" + metric_log = cfg.metric_logfile or f"pipeline_metric_{os.getpid()}.log" + action_logger = _setup_logger("pipeline_action_logger", action_log, level=logging.DEBUG) + metric_logger = _setup_logger("pipeline_metric_logger", metric_log) + + for flag in ("webhook", "checkpoint_on_failure", "cgroup"): + if getattr(ns, flag, None): + action_logger.warning("--%s is accepted but has no effect", + flag.replace("_", "-")) + + # also route the package-level log records to the action log + pkg_log = logging.getLogger("o2dpg_runner") + pkg_log.setLevel(logging.INFO) + for h in list(pkg_log.handlers): + pkg_log.removeHandler(h) + pkg_log.propagate = False + for h in action_logger.handlers: + pkg_log.addHandler(h) + + # Apply slice-level cgroup resource limits now that we are inside the + # slice and the action logger is ready to record the outcome. + if cfg.in_systemd_slice and cfg.systemd_run_spec: + try: + cpu_quota, mem_str, _ = _parse_systemd_run_spec(cfg.systemd_run_spec) + _apply_slice_cgroup_limits(cpu_quota, mem_str, action_logger) + except Exception as e: + action_logger.warning("Could not apply slice cgroup limits: %s", e) + + # record meta to the metric log (mirrors prototype) + raw = load_json(cfg.workflowfile) + meta = raw.get("meta", {}) if isinstance(raw, dict) else {} + if not isinstance(meta, dict): + meta = {} + meta.update({ + "cpu_limit": cfg.cpu_limit, + "mem_limit": cfg.mem_limit, + "workflow_file": os.path.abspath(cfg.workflowfile), + "target_task": cfg.target_tasks, + "rerun_from": cfg.rerun_from, + "target_labels": cfg.target_labels, + "scheduler_policy": cfg.scheduler_policy, + "drop_should_break": cfg.drop_should_break, + "cache_policy": cfg.cache_policy, + "systemd_run_spec": cfg.systemd_run_spec, + "in_systemd_slice": cfg.in_systemd_slice, + "monitor_interval_cpu": cfg.monitor_interval_cpu, + }) + metric_logger.info(meta) + + # visualize if asked (uses raw spec before filtering) + if cfg.visualize_workflow: + _maybe_draw_workflow(raw) + + # build workflow (filters, strips global init, builds DAG) + wf = build_workflow(raw, cfg.target_tasks, cfg.target_labels) + if not wf.stages: + if cfg.target_tasks: + print("Apparently some of the chosen target tasks are not in the workflow") + else: + print("Workflow is empty. Nothing to do") + return 0 + + # Apply global env (as the prototype did at construction time) + for k, v in wf.global_env.items(): + os.environ.setdefault(k, str(v)) + + # Optional file-access sidecar + fileaccess_proc, fileaccess_fh, fileaccess_log_file = _launch_fileaccess_sidecar(action_log) + + rc = 0 + try: + execer = WorkflowExecutor(cfg, wf, action_logger, metric_logger) + rc = int(execer.execute()) + finally: + if fileaccess_proc is not None: + fileaccess_proc.terminate() + try: + fileaccess_proc.wait(timeout=5) + except subprocess.TimeoutExpired: + fileaccess_proc.kill() + if fileaccess_fh is not None: + fileaccess_fh.close() + o2dpg_root = os.getenv("O2DPG_ROOT") + if o2dpg_root and fileaccess_log_file: + analyse_cmd = [ + sys.executable, + f"{o2dpg_root}/UTILS/FileIOGraph/analyse_FileIO_v2.py", + "--actionFile", action_log, + "--monitorFile", fileaccess_log_file, + "-o", f"pipeline_fileaccess_report_{os.getpid()}.json", + "--basedir", os.getcwd(), + ] + print(f"Producing FileIOGraph with command {analyse_cmd}") + try: + subprocess.run(analyse_cmd, check=True) + except subprocess.CalledProcessError as e: + print(f"FileIOGraph analysis failed: {e}", file=sys.stderr) + + return rc + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/MC/workflow_runner/o2dpg_runner/config.py b/MC/workflow_runner/o2dpg_runner/config.py new file mode 100644 index 000000000..d12bf73c8 --- /dev/null +++ b/MC/workflow_runner/o2dpg_runner/config.py @@ -0,0 +1,62 @@ +"""Runtime configuration for the runner. + +A dataclass that holds everything the old module-level ``args`` exposed, +so nothing in the code needs to import argparse. Constructed from the +argparse Namespace in cli.py. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import List, Optional + + +@dataclass +class RunnerConfig: + # --- required --- + workflowfile: str + + # --- scheduling / resources --- + maxjobs: int = 100 + mem_limit: float = 0.0 # MB; 0 means "auto from psutil" + cpu_limit: float = 8.0 + n_backfill: int = 1 + update_resources: Optional[str] = None + dynamic_resources: bool = False + optimistic_resources: bool = False + in_systemd_slice: bool = False # True when runner was re-exec'd under systemd-run --scope + + # --- new scheduler knobs --- + scheduler_policy: str = "timeframe" # timeframe | critical-path | best-fit + drop_should_break: bool = False # let timeframe policy scan past non-fitting + + # --- systemd-run slice --- + systemd_run_spec: Optional[str] = None # raw "ncpus:N/mem:M/name:S" spec, kept for metric meta + systemd_slice_name: Optional[str] = None # parsed "name:" value; used for child scope names + + # --- new monitor knobs --- + monitor_interval_cpu: float = 1.0 + monitor_interval_mem: float = 1.0 # match prototype cadence; raise for cheaper monitor + monitor_backend: str = "psutil" # psutil | auto (auto reserved for future cgroup backend) + + # --- cache policy (v1 _done.json) --- + cache_policy: str = "off" # off | lenient | strict + + # --- selection / control --- + target_tasks: List[str] = field(default_factory=lambda: ["*"]) + target_labels: List[str] = field(default_factory=list) + keep_going: bool = False + dry_run: bool = False + visualize_workflow: bool = False + produce_script: Optional[str] = None + rerun_from: Optional[str] = None + list_tasks: bool = False + retry_on_failure: int = 0 + no_rootinit_speedup: bool = False + remove_files_early: str = "" + stdout_on_failure: bool = False + production_mode: bool = False + + # --- logging --- + action_logfile: Optional[str] = None + metric_logfile: Optional[str] = None diff --git a/MC/workflow_runner/o2dpg_runner/executor.py b/MC/workflow_runner/o2dpg_runner/executor.py new file mode 100644 index 000000000..10dc73b45 --- /dev/null +++ b/MC/workflow_runner/o2dpg_runner/executor.py @@ -0,0 +1,763 @@ +"""Main control loop. + +Responsibilities: + - boot sequence (global init, ROOT speedup, FAIRMQ socket setup) + - candidate management + - submit / wait / monitor-feedback interplay + - retry, failure handling, rerun-from + - end-of-task hooks (file removal, production archival) + +Split from the prototype's monolithic WorkflowExecutor and freed from +module-level args. +""" + +from __future__ import annotations + +import json +import logging +import os +import platform +import re +import signal +import subprocess +import sys +import threading +import time +import traceback +from dataclasses import dataclass, field +from typing import Dict, List, Optional, Set, Tuple + +import psutil + +from .config import RunnerConfig +from .workflow import Workflow, update_resource_estimates +from .graph import descendants, longest_path_length, kahn_topological_order +from .resources import ResourceManager, ResourceLimitExceeded +from .monitoring import MonitorThread, PsutilBackend, _read_cgroup_v2_dir +from .scheduler import get_policy +from .scheduler.base import SchedulerState +from .scheduler.timeframe import TimeframeFirstPolicy +from .cache import TaskCache, compute_fingerprint, remove_done_flag, done_path +from .alienv import get_alienv_software_environment +from .cleanup import EarlyFileRemover, archive_task_logs + +log = logging.getLogger(__name__) + +_UNIT_NAME_RE = re.compile(r"[^a-zA-Z0-9_\-.]") + + +def _unit_name(task_name: str, tid: int) -> str: + """Build a valid systemd unit name for a per-task scope.""" + safe = _UNIT_NAME_RE.sub("-", task_name) + return f"task-{safe}-{tid}.scope" + + +def _start_stderr_drainer(pipe, logger: logging.Logger, tag: str) -> threading.Thread: + """Drain *pipe* line-by-line in a daemon thread, forwarding to *logger*. + + Used to capture systemd-run's own informational messages (e.g. "Running + as unit: ...") and route them to the action log instead of the terminal. + The thread exits naturally when the pipe reaches EOF (process finished). + """ + def _run() -> None: + try: + for raw in pipe: + line = raw.rstrip() if isinstance(raw, str) else raw.rstrip().decode(errors="replace") + if line: + logger.info("[systemd] %s: %s", tag, line) + except Exception: + pass + t = threading.Thread(target=_run, daemon=True, name=f"stderr-{tag}") + t.start() + return t + + +@dataclass +class _TaskRuntime: + """Everything we learn about a task at runtime and need to feed back.""" + logfile: str + fingerprint: Dict[str, str] = field(default_factory=dict) + start_time: float = 0.0 + pid: int = 0 + + +class WorkflowExecutor: + """Construct from a config + workflow; call execute().""" + + def __init__( + self, + config: RunnerConfig, + workflow: Workflow, + action_logger: logging.Logger, + metric_logger: logging.Logger, + ): + self.cfg = config + self.wf = workflow + self.actionlog = action_logger + self.metriclog = metric_logger + + # apply update-resources (before building resource manager) + if config.update_resources: + update_resource_estimates( + workflow, config.update_resources, + logger=action_logger, + ) + + # resource manager + self.rm = ResourceManager( + cpu_limit=config.cpu_limit, + mem_limit=config.mem_limit, + procs_parallel_max=config.maxjobs, + n_backfill_max=config.n_backfill, + dynamic_resources=config.dynamic_resources, + optimistic_resources=config.optimistic_resources, + ) + for task in workflow.stages: + try: + rel = float(task["resources"].get("relative_cpu") or 1) + except (TypeError, ValueError): + rel = 1.0 + try: + self.rm.add_task( + name=task["name"], + related_name=self._global_name(task["name"]), + cpu=float(task["resources"]["cpu"]), + cpu_relative=rel, + mem=float(task["resources"]["mem"]), + semaphore_string=task.get("semaphore"), + ) + except ResourceLimitExceeded as e: + print(e, file=sys.stderr) + print("Pass --optimistic-resources to the runner to attempt the run anyway.", + file=sys.stderr) + raise + + # scheduler + if config.scheduler_policy == "timeframe": + self.policy = TimeframeFirstPolicy(drop_should_break=config.drop_should_break) + else: + self.policy = get_policy(config.scheduler_policy) + + # scheduler state: precompute weights once + self.state = self._build_scheduler_state() + + # task cache (covers _done + optional _done.json) + self.cache = TaskCache(policy=config.cache_policy) + # precompute per-task fingerprint so we don't recompute on each check + self._fingerprint_by_tid: Dict[int, Dict[str, str]] = {} + for tid, task in enumerate(workflow.stages): + alienv = task.get("alternative_alienv_package") or "" + self._fingerprint_by_tid[tid] = compute_fingerprint(task, alienv) + + # alternative alienv envs + self.alternative_envs: Dict[int, Dict[str, str]] = {} + self._init_alternative_envs() + + # Compute the global cgroup directory for the aggregate monitor. + # Cgroup monitoring is only meaningful when the runner was launched + # under --systemd-run: the runner's own scope is a leaf node + # (e.g. o2dpg.slice/o2dpg-runner-.scope/) and task scopes are + # siblings. The *parent* slice directory's cpu.stat / memory.current + # then cover the runner + all task scopes. Without --systemd-run the + # runner shares a generic user-session cgroup with unrelated + # processes, so cgroup readings would be misleading; fall back to + # plain psutil monitoring instead. + _global_cgroup: Optional[str] = None + if config.in_systemd_slice: + _runner_cgroup = _read_cgroup_v2_dir(os.getpid()) + if _runner_cgroup: + _global_cgroup = os.path.dirname(_runner_cgroup) + if not os.path.isdir(_global_cgroup): + _global_cgroup = _runner_cgroup # safety fallback + + # monitor + self.monitor = MonitorThread( + cpu_interval=config.monitor_interval_cpu, + mem_interval=config.monitor_interval_mem, + backend=PsutilBackend(), + monitor_disc=bool(os.getenv("MONITOR_DISC_USAGE")), + disc_path=os.getcwd(), + global_cgroup_dir=_global_cgroup, + ) + + # process tracking + self.proc_status: Dict[int, str] = {tid: "ToDo" for tid in range(workflow.n_tasks())} + self.task_runtime: Dict[int, _TaskRuntime] = {} + self.process_list: List[Tuple[int, psutil.Popen]] = [] + self.tids_marked_retry: List[int] = [] + self.retry_counter: List[int] = [0] * workflow.n_tasks() + self.task_retries: List[int] = [ + int(t.get("retry_count", 0)) for t in workflow.stages + ] + + # early file removal + self.file_remover: Optional[EarlyFileRemover] = None + if config.remove_files_early: + try: + self.file_remover = EarlyFileRemover( + config.remove_files_early, + workflow.timeframes, + workflow.full_target_names, + logger=action_logger, + ) + except Exception as e: + log.warning("Could not set up early file removal: %s", e) + + self.start_time: float = 0.0 + self.scheduling_iteration = 0 + self._last_metric_tick: int = -1 # prevents duplicate metric rows per tick + + # signals + signal.signal(signal.SIGINT, self._sighandler) + signal.signal(signal.SIGTERM, self._sighandler) + signal.siginterrupt(signal.SIGINT, False) + signal.siginterrupt(signal.SIGTERM, False) + + # ----- small helpers ----- + @staticmethod + def _global_name(name: str) -> str: + """Strip _ suffix to find sibling group for resource sampling.""" + toks = name.split("_") + if toks and toks[-1].isdigit() and len(toks) > 1: + return "_".join(toks[:-1]) + return name + + def _build_scheduler_state(self) -> SchedulerState: + n = self.wf.n_tasks() + # descendants via memoized iterative DFS (no recursion limit concerns) + desc_cache: Dict[int, Set[int]] = {} + desc_counts = [0] * n + for tid in range(n): + desc_counts[tid] = len(descendants(self.wf.forward_adj, tid, desc_cache)) + + timeframe_of = [t.get("timeframe", -1) for t in self.wf.stages] + tf_weight = [(timeframe_of[t], desc_counts[t]) for t in range(n)] + + cpu = [float(t.get("resources", {}).get("cpu", 1.0)) for t in self.wf.stages] + mem = [float(t.get("resources", {}).get("mem", 0.0)) for t in self.wf.stages] + + # Per-task walltime [s] from learned resources (resources.walltime set + # by update_resource_estimates when --update-resources is given). + # Fall back to cpu as a proxy so behaviour is unchanged without + # learned data. + walltime = [ + float(t.get("resources", {}).get("walltime") or cpu[i]) + for i, t in enumerate(self.wf.stages) + ] + has_walltime = any( + t.get("resources", {}).get("walltime") for t in self.wf.stages + ) + + # Critical path: longest remaining *wall time* to any leaf. + # Using walltime as the node weight gives a true makespan estimate; + # using cpu (the fallback) preserves the original heuristic. + cp_weight = walltime if has_walltime else cpu + topo = kahn_topological_order(n, self.wf.forward_adj, self.wf.indegree) + cp = longest_path_length(self.wf.forward_adj, topo, cp_weight) + + if has_walltime: + self.actionlog.info("Critical path weighted by learned walltime [s]") + + # self-log weights (matches prototype's informational logging) + for tid in range(n): + self.actionlog.info("Score for %s is %s", self.wf.id_to_name[tid], tf_weight[tid]) + + return SchedulerState( + timeframe_of=timeframe_of, + descendants_count=desc_counts, + critical_path=cp, + task_cpu=cpu, + task_mem=mem, + task_walltime=walltime, + timeframe_weight=tf_weight, + ) + + def _init_alternative_envs(self) -> None: + cache: Dict[str, Dict[str, str]] = {} + for tid, task in enumerate(self.wf.stages): + pkg = task.get("alternative_alienv_package") + if not pkg: + continue + if pkg not in cache: + cache[pkg] = get_alienv_software_environment(pkg) + self.alternative_envs[tid] = cache[pkg] + + # ----- task-level helpers ----- + def logfile(self, tid: int) -> str: + task = self.wf.stages[tid] + return os.path.join(task.get("cwd", "."), f"{task['name']}.log") + + def apply_global_env(self, env: Dict[str, str]) -> None: + for k, v in self.wf.global_env.items(): + env.setdefault(k, str(v)) + + # ----- signal handling ----- + def _sighandler(self, signum, frame): + self.actionlog.info("Signal %s caught; terminating children", signum) + try: + self.monitor.stop() + except Exception: + pass + try: + procs = psutil.Process().children(recursive=True) + except psutil.NoSuchProcess: + procs = [] + except (psutil.AccessDenied, PermissionError): + procs = [] + + for p in procs: + try: + p.terminate() + except (psutil.NoSuchProcess, psutil.AccessDenied): + pass + _, alive = psutil.wait_procs(procs, timeout=3) + for p in alive: + try: + p.kill() + except (psutil.NoSuchProcess, psutil.AccessDenied): + pass + sys.exit(1) + + # ----- task submission ----- + def submit(self, tid: int, nice: int) -> Optional[psutil.Popen]: + task = self.wf.stages[tid] + self.actionlog.debug("Submitting %s with nice=%d", task["name"], nice) + cmd = task["cmd"] + workdir = task.get("cwd", ".") + if workdir: + if os.path.exists(workdir) and not os.path.isdir(workdir): + self.actionlog.error("cwd %s exists and is not a directory", workdir) + return None + if not os.path.isdir(workdir): + os.makedirs(workdir, exist_ok=True) + + self.proc_status[tid] = "Running" + + if self.cfg.dry_run: + dry = f"echo ' {self.scheduling_iteration} : would do {task['name']}'" + # psutil.Popen so that the dry-run path also answers .nice() + return psutil.Popen(["/bin/bash", "-c", dry], cwd=workdir) + + env = os.environ.copy() + alt = self.alternative_envs.get(tid) + if alt: + self.actionlog.info("Applying alternative environment to %s", task["name"]) + if alt.get("TERM") is not None: + env = dict(alt) + else: + env.update(alt) + if task.get("env"): + env.update({k: str(v) for k, v in task["env"].items()}) + self.apply_global_env(env) + + if os.environ.get("PIPELINE_RUNNER_DUMP_TASKENVS") is not None: + try: + with open(f"taskenv_{tid}.log", "w") as f: + json.dump(env, f, indent=2) + except OSError as e: + log.warning("could not dump taskenv: %s", e) + + # When the runner is inside a systemd slice, wrap each task in its own + # child scope so per-task cgroup metrics are available alongside psutil. + slice_name = self.cfg.systemd_slice_name + use_scope = self.cfg.in_systemd_slice and bool(slice_name) + if use_scope: + systemd_slice = ( + slice_name if slice_name.endswith(".slice") else f"{slice_name}.slice" + ) + unit = _unit_name(task["name"], tid) + launch_argv = [ + "systemd-run", "--user", "--scope", "--collect", + "--expand-environment=no", # suppress the $VAR warning; bash handles expansion + f"--unit={unit}", f"--slice={systemd_slice}", + "--", "/bin/bash", "-c", cmd, + ] + p = psutil.Popen(launch_argv, cwd=workdir, env=env, stderr=subprocess.PIPE) + _start_stderr_drainer(p.stderr, self.actionlog, task["name"]) + else: + launch_argv = ["/bin/bash", "-c", cmd] + p = psutil.Popen(launch_argv, cwd=workdir, env=env) + try: + p.nice(nice) + except (psutil.NoSuchProcess, psutil.AccessDenied): + self.actionlog.error("Could not renice %d to %d", p.pid, nice) + + rt = _TaskRuntime( + logfile=self.logfile(tid), + fingerprint=self._fingerprint_by_tid[tid], + start_time=time.perf_counter(), + pid=p.pid, + ) + self.task_runtime[tid] = rt + self.monitor.register( + tid, p.pid, task["name"], task.get("labels", []) or [], rt.start_time, + resolve_cgroup=use_scope, + ) + return p + + # ----- skip logic ----- + def ok_to_skip(self, tid: int) -> bool: + return self.cache.is_done(self.logfile(tid), self._fingerprint_by_tid[tid]) + + # ----- candidate scheduling pass ----- + def try_submit_from_candidates( + self, + candidates: List[int], + finished_out: List[int], + ) -> None: + self.scheduling_iteration += 1 + + # skip already-done tasks first (in place) + remaining = [] + for tid in candidates: + if self.ok_to_skip(tid): + finished_out.append(tid) + self.actionlog.info("Skipping %s", self.wf.id_to_name[tid]) + if self.file_remover is not None: + self.file_remover.on_task_done(self.wf.id_to_name[tid]) + else: + remaining.append(tid) + # mutate the list the caller passed in + candidates[:] = remaining + + ordered = self.policy.order(candidates, self.state) + for tid, nice in self.policy.pick_submittable(ordered, self.rm): + self.actionlog.debug("Submitting tid=%d %s (nice=%d)", + tid, self.wf.id_to_name[tid], nice) + p = self.submit(tid, nice) + if p is None: + continue + # pin the nice value the OS actually granted + try: + actual_nice = p.nice() + except (psutil.NoSuchProcess, psutil.AccessDenied): + actual_nice = nice + self.rm.book(tid, actual_nice) + self.process_list.append((tid, p)) + if tid in candidates: + candidates.remove(tid) + + # ----- wait / complete / retry ----- + def wait_for_any( + self, + finished_out: List[int], + failing_out: List[int], + ) -> bool: + """Return True if we should keep waiting (no completion this pass). + + Polls each process via poll(); pulls latest monitor snapshot and + feeds it into the ResourceManager for the dynamic-resources path. + """ + if not self.process_list: + return False + + # Take each monitor tick exactly once. This loop polls faster than the + # monitor fires, and a snapshot recorded twice both inflates the sample + # lists and defeats the "too few samples" guard in sample_resources(). + snapshots = self.monitor.latest() + tick = self.monitor.tick + + if tick != self._last_metric_tick: + self._last_metric_tick = tick + for tid, snap in snapshots.items(): + self.rm.add_monitored(tid, snap.t_delta_ms, + snap.cpu_pct / 100.0, snap.pss_mb) + self.metriclog.info({ + "iter": tick, "name": snap.name, + "cpu": snap.cpu_pct, "uss": snap.uss_mb, "pss": snap.pss_mb, + "nice": snap.nice, "swap": snap.swap_mb, + "label": snap.labels, "disc": snap.disc_mb, + # cgroup-based readings for comparison with psutil (None when + # no per-task scope is active) + "cgroup_cpu": snap.cgroup_cpu_pct, "cgroup_mem": snap.cgroup_mem_mb, + }) + + # cgroup-aggregate slice totals + g_cpu = self.monitor.global_cpu_pct + g_mem = self.monitor.global_mem_mb + if g_cpu is not None or g_mem is not None: + self.metriclog.info({ + "iter": tick, "name": "__cgroup_global__", + "cpu": g_cpu, "uss": None, "pss": g_mem, + "nice": 0, "swap": None, "label": [], "disc": -1, + }) + + # check for completions + newly_done: List[Tuple[int, psutil.Popen, int]] = [] + for tid, p in list(self.process_list): + rc = 0 if self.cfg.dry_run else p.poll() + if rc is None: + continue + newly_done.append((tid, p, rc)) + + failure_detected = False + for tid, p, rc in newly_done: + name = self.wf.id_to_name[tid] + self.actionlog.info("Task pid=%d tid=%d %s finished rc=%d", + p.pid, tid, name, rc) + self.rm.unbook(tid) + self.proc_status[tid] = "Done" + self.monitor.deregister(tid) + self.process_list.remove((tid, p)) + + if rc == 0: + finished_out.append(tid) + # record fingerprint sidecar (best effort) + rt = self.task_runtime.get(tid) + if rt is not None: + self.cache.record(rt.logfile, rt.fingerprint) + if self.file_remover is not None: + self.file_remover.on_task_done(name) + if self.cfg.production_mode: + archive_task_logs(self.logfile(tid), logger=self.actionlog) + else: + print(f"{name} failed ... checking retry") + max_retries = max(self.cfg.retry_on_failure, self.task_retries[tid]) + if self._is_worth_retrying(tid) and self.retry_counter[tid] < max_retries: + self.actionlog.info("Task %s marked for retry", name) + self.tids_marked_retry.append(tid) + self.retry_counter[tid] += 1 + else: + failure_detected = True + failing_out.append(tid) + + if failure_detected and not self.cfg.keep_going: + self.actionlog.info("Stopping due to failure in tids %s", failing_out) + if self.cfg.stdout_on_failure: + self._cat_logfiles(failing_out) + self.stop_and_exit() + + return not finished_out + + def _is_worth_retrying(self, tid: int) -> bool: + """Hook for future log-inspection; currently always True (match prototype).""" + return True + + def _cat_logfiles(self, tids: List[int]) -> None: + for tid in tids: + logf = self.logfile(tid) + if os.path.exists(logf): + print(f" ----> START OF LOGFILE {logf} -----") + try: + with open(logf) as f: + sys.stdout.write(f.read()) + except OSError: + pass + print(f" <---- END OF LOGFILE {logf} -----") + + def stop_and_exit(self) -> None: + for _, p in self.process_list: + try: + p.kill() + except Exception: + pass + self.monitor.stop() + sys.exit(1) + + # ----- boot helpers ----- + def _speedup_root_init(self) -> None: + if platform.system() != "Linux": + return + if os.environ.get("ROOT_LDSYSPATH") and os.environ.get("ROOT_CPPSYSINCL"): + return + if self.cfg.no_rootinit_speedup: + return + try: + cmd = ('LD_DEBUG=libs LD_PRELOAD=DOESNOTEXIST ls /tmp/DOESNOTEXIST 2>&1 | ' + 'grep -m 1 "system search path" | sed \'s/.*=//g\' | ' + 'awk \'//{print $1}\'') + libpath = subprocess.check_output(cmd, shell=True).decode().strip() + if libpath: + os.environ["ROOT_LDSYSPATH"] = libpath + os.environ["CLING_LDSYSPATH"] = libpath + cmd2 = ("LC_ALL=C c++ -xc++ -E -v /dev/null 2>&1 | " + "sed -n '/^#include/,${/^ \\/.*++/{p}}'") + incpath = subprocess.check_output(cmd2, shell=True).decode() + joined = ":".join(line.lstrip() for line in incpath.splitlines()) + if joined: + os.environ["ROOT_CPPSYSINCL"] = joined + os.environ["CLING_CPPSYSINCL"] = joined + except Exception as e: + log.warning("ROOT init speedup failed: %s", e) + + def _execute_global_init_cmd(self) -> bool: + cmd = self.wf.global_init_cmd + if not cmd: + return True + self.actionlog.info("Executing global init cmd: %s", cmd) + p = subprocess.Popen(["/bin/bash", "-c", cmd], + stdout=subprocess.PIPE, stderr=subprocess.PIPE) + stdout, stderr = p.communicate() + if p.returncode == 0: + self.actionlog.info(stdout.decode()) + return True + self.actionlog.error("global init failed: %s", stderr.decode()) + return False + + def _handle_rerun_from(self) -> None: + if not self.cfg.rerun_from: + return + import re + matched = False + for task in self.wf.stages: + if re.match(self.cfg.rerun_from, task["name"]): + matched = True + tid = self.wf.tid(task["name"]) + # remove done flags for tid and all its descendants (iterative) + for d in descendants(self.wf.forward_adj, tid) | {tid}: + name = self.wf.id_to_name[d] + self.actionlog.info("Marking %s for rerun", name) + if not self.cfg.dry_run: + remove_done_flag(self.logfile(d)) + else: + print(f"Would mark {name} as to be done again") + if not matched: + print(f"No task matching {self.cfg.rerun_from} found; refusing to proceed") + sys.exit(1) + + # ----- bash-script emission ----- + def produce_script(self, filename: str) -> None: + topo = kahn_topological_order(self.wf.n_tasks(), + self.wf.forward_adj, + self.wf.indegree) + lines = [ + "#!/usr/bin/env bash\n", + "#THIS FILE IS AUTOGENERATED\n", + "export JOBUTILS_SKIPDONE=ON\n", + "#-- GLOBAL INIT SECTION FROM WORKFLOW --\n", + ] + for k, v in self.wf.global_env.items(): + lines.append(f"export {k}={v}\n") + lines.append("#-- TASKS FROM WORKFLOW --\n") + for tid in topo: + t = self.wf.stages[tid] + workdir = t.get("cwd", ".") + env_pairs = t.get("env") or {} + env_prefix = " ".join(f"{k}={v}" for k, v in env_pairs.items()) + # Subshell so inner `cd` doesn't leak, and local env doesn't pollute. + inner = t["cmd"] + if env_prefix: + inner = f"{env_prefix} {inner}" + lines.append(f"( [ -d {workdir} ] || mkdir -p {workdir}; cd {workdir} && {inner} )\n") + with open(filename, "w") as f: + f.writelines(lines) + + # ----- main loop ----- + def execute(self) -> bool: + self.start_time = time.perf_counter() + psutil.cpu_percent(interval=None) + os.environ["JOBUTILS_SKIPDONE"] = "ON" + self._speedup_root_init() + + if not os.path.isdir("./.tmp"): + os.mkdir("./.tmp") + if os.environ.get("FAIRMQ_IPC_PREFIX") is None: + sp = os.path.join(os.getcwd(), ".tmp") + self.actionlog.info("Setting FAIRMQ_IPC_PREFIX=%s", sp) + os.environ["FAIRMQ_IPC_PREFIX"] = sp + + if self.cfg.list_tasks: + print("List of tasks in this workflow:") + for i, t in enumerate(self.wf.stages): + label_part = t.get("labels", []) + print(f"{t['name']} ({label_part}) ToDo: {not self.ok_to_skip(i)}") + return False + + if self.cfg.produce_script is not None: + self.produce_script(self.cfg.produce_script) + return False + + if not self._execute_global_init_cmd(): + sys.exit(1) + + self._handle_rerun_from() + + # start monitor + self.monitor.start() + + # initial candidates: tasks with no predecessors + candidates = [i for i, d in enumerate(self.wf.indegree) if d == 0] + finishedtasks_set: Set[int] = set() + error_encountered = False + + try: + while True: + finished: List[int] = [] + self.actionlog.debug("candidates: %s", + [(c, self.wf.id_to_name[c]) for c in candidates]) + self.try_submit_from_candidates(candidates, finished) + + if candidates and not self.process_list: + self._noprogress_error() + error_encountered = True + break + + # wait loop + finished_running: List[int] = [] + failing: List[int] = [] + poll_delay = 0.1 # adaptive; grows up to 1s + while self.wait_for_any(finished_running, failing): + if not self.cfg.dry_run: + time.sleep(poll_delay) + poll_delay = min(1.0, poll_delay * 1.5) + else: + time.sleep(0.001) + + finished.extend(finished_running) + finishedtasks_set.update(finished) + + # take failed tasks out of the "finished" accounting + if failing: + error_encountered = True + fs = set(failing) + finished = [x for x in finished if x not in fs] + finishedtasks_set.difference_update(fs) + + # retries go back onto the candidate list + if self.tids_marked_retry: + rs = set(self.tids_marked_retry) + finished = [x for x in finished if x not in rs] + finishedtasks_set.difference_update(rs) + for t in self.tids_marked_retry: + if t not in candidates: + candidates.append(t) + self.tids_marked_retry.clear() + + # new candidates: successors whose all needs are done + for tid in finished: + for succ in self.wf.forward_adj[tid]: + if succ in candidates: + continue + if self.proc_status[succ] != "ToDo": + continue + preds = self.wf.reverse_adj[succ] + if all(p in finishedtasks_set for p in preds): + candidates.append(succ) + + self.actionlog.debug("new candidates %s", candidates) + + if not candidates and not self.process_list: + break + except Exception: + traceback.print_exc() + self._sighandler(0, None) + + self.monitor.stop() + self.monitor.join(timeout=2) + end = time.perf_counter() + msg = "with failures" if error_encountered else "success" + print(f"\n**** Pipeline done {msg} (global_runtime : {end - self.start_time:.3f}s) *****\n") + self.actionlog.debug("global_runtime : %.3fs", end - self.start_time) + return error_encountered + + # ----- error message ----- + def _noprogress_error(self) -> None: + msg = ( + "Scheduler runtime error: cannot make progress although candidates exist.\n\n" + "This typically means a task's estimated resources exceed the configured\n" + "--cpu-limit or --mem-limit. On a 16 GB node, try --mem-limit 20000 (MB);\n" + "the ACTUAL use might be lower than the estimate. Alternatively convert\n" + "the workflow to a linear shell script via --produce-script .sh and\n" + "run that directly.\n" + ) + print(msg, file=sys.stderr) diff --git a/MC/workflow_runner/o2dpg_runner/graph.py b/MC/workflow_runner/o2dpg_runner/graph.py new file mode 100644 index 000000000..d04184f47 --- /dev/null +++ b/MC/workflow_runner/o2dpg_runner/graph.py @@ -0,0 +1,159 @@ +"""Graph utilities on the task DAG. + +Nothing here knows about workflows, tasks, or resources -- it operates on +integer-indexed adjacency lists only. Pure, testable, no side effects. + +Supersedes the recursive ``findAllTopologicalOrders`` and +``find_all_dependent_tasks`` in the original prototype, both of which had +bugs and needed sys.setrecursionlimit(100000) to survive diamond DAGs. +""" + +from __future__ import annotations + +from collections import defaultdict, deque +from typing import Dict, Iterable, List, Set, Tuple + + +def build_adjacency( + n_nodes: int, edges: Iterable[Tuple[int, int]] +) -> Tuple[List[List[int]], List[List[int]], List[int]]: + """Return (forward_adj, reverse_adj, indegree) for ``n_nodes`` nodes. + + forward_adj[u] lists successors; reverse_adj[v] lists predecessors. + """ + forward: List[List[int]] = [[] for _ in range(n_nodes)] + reverse: List[List[int]] = [[] for _ in range(n_nodes)] + indeg = [0] * n_nodes + for u, v in edges: + forward[u].append(v) + reverse[v].append(u) + indeg[v] += 1 + return forward, reverse, indeg + + +def kahn_topological_order( + n_nodes: int, + forward_adj: List[List[int]], + indegree: List[int], + tiebreak: List[int] = None, +) -> List[int]: + """Deterministic topological order via Kahn's algorithm. + + ``tiebreak`` is an optional per-node integer used to break ties + among ready nodes. Smaller tiebreak value first. If None, the + node index is used (which is also deterministic). + """ + if tiebreak is None: + tiebreak = list(range(n_nodes)) + indeg = list(indegree) + # Use a sorted list as a tiny priority queue; for small DAGs this is + # fine, and ties are broken deterministically. + ready = sorted([n for n in range(n_nodes) if indeg[n] == 0], + key=lambda n: (tiebreak[n], n)) + out: List[int] = [] + # Simple loop; no heap because re-sorting on small ready sets is cheap + # and we want full determinism. + while ready: + u = ready.pop(0) + out.append(u) + for v in forward_adj[u]: + indeg[v] -= 1 + if indeg[v] == 0: + # insert keeping sorted-by-tiebreak + lo, hi = 0, len(ready) + key = (tiebreak[v], v) + while lo < hi: + mid = (lo + hi) // 2 + if (tiebreak[ready[mid]], ready[mid]) < key: + lo = mid + 1 + else: + hi = mid + ready.insert(lo, v) + if len(out) != n_nodes: + raise ValueError("Graph has at least one cycle; topological sort impossible") + return out + + +def descendants( + forward_adj: List[List[int]], source: int, cache: Dict[int, Set[int]] = None +) -> Set[int]: + """All nodes reachable from ``source`` (excluding ``source`` itself), memoized. + + Uses iterative DFS with post-order memoization. Safe on DAGs with + diamonds and deep chains; no recursion limit concerns. + """ + if cache is None: + cache = {} + if source in cache: + return cache[source] + + # iterative post-order: process children first, then combine + order: List[int] = [] + seen: Set[int] = set() + stack: List[Tuple[int, int]] = [(source, 0)] + while stack: + node, child_idx = stack[-1] + kids = forward_adj[node] + if child_idx < len(kids): + stack[-1] = (node, child_idx + 1) + child = kids[child_idx] + if child not in seen and child not in cache: + seen.add(child) + stack.append((child, 0)) + else: + stack.pop() + order.append(node) + + # Now assign cache entries in post-order (children before parents). + for node in order: + if node in cache: + continue + s: Set[int] = set() + for child in forward_adj[node]: + s.add(child) + s |= cache.get(child, set()) + cache[node] = s + return cache[source] + + +def ancestors( + reverse_adj: List[List[int]], sink: int, cache: Dict[int, Set[int]] = None +) -> Set[int]: + """All nodes that can reach ``sink`` (excluding sink itself), memoized.""" + # ancestors in forward graph == descendants in reverse graph + return descendants(reverse_adj, sink, cache) + + +def longest_path_length( + forward_adj: List[List[int]], + topo_order: List[int], + node_weight: List[float], +) -> List[float]: + """Longest-path weight from each node to any leaf (inclusive of node). + + ``topo_order`` must be a valid topological ordering. Uses reverse + traversal and DP. Used by CriticalPathPolicy. + """ + n = len(node_weight) + lp = list(node_weight) # include own weight + for u in reversed(topo_order): + best_child = 0.0 + for v in forward_adj[u]: + if lp[v] > best_child: + best_child = lp[v] + lp[u] = node_weight[u] + best_child + return lp + + +def invert_adj(forward_adj: List[List[int]]) -> List[List[int]]: + """Compute reverse adjacency from forward adjacency.""" + n = len(forward_adj) + rev: List[List[int]] = [[] for _ in range(n)] + for u in range(n): + for v in forward_adj[u]: + rev[v].append(u) + return rev + + +def root_nodes(indegree: List[int]) -> List[int]: + return [i for i, d in enumerate(indegree) if d == 0] diff --git a/MC/workflow_runner/o2dpg_runner/monitoring.py b/MC/workflow_runner/o2dpg_runner/monitoring.py new file mode 100644 index 000000000..f156c6655 --- /dev/null +++ b/MC/workflow_runner/o2dpg_runner/monitoring.py @@ -0,0 +1,505 @@ +"""Resource monitor running in a background thread. + +The main scheduler loop used to poll psutil synchronously at 1 Hz, which +cost 10-20% of one core on realistic workflows. Now the monitor is a +separate thread with independent CPU (cheap, 1 Hz) and MEM (expensive, +0.2 Hz) cadences. The scheduler reads the latest snapshot lock-free via +an atomic dict reference. + +Backend interface: a callable + sample(task_pid_list) -> {pid: {"cpu_pct": float, "pss_mb": float, + "uss_mb": float, "swap_mb": float, + "children": [pid, ...]}} +""" + +from __future__ import annotations + +import logging +import os +import subprocess +import threading +import time +from dataclasses import dataclass, field +from typing import Dict, List, Optional, Tuple + +try: + import psutil + HAVE_PSUTIL = True +except ImportError: + HAVE_PSUTIL = False + +log = logging.getLogger(__name__) + + +@dataclass +class TaskSnapshot: + """Per-task roll-up of the most recent monitor pass.""" + tid: int + name: str + t_delta_ms: int = 0 + cpu_pct: float = 0.0 # 0..100 * n_cores (psutil) + uss_mb: float = 0.0 + pss_mb: float = 0.0 + swap_mb: float = 0.0 + nice: int = 0 + labels: List[str] = field(default_factory=list) + disc_mb: float = -1.0 # global disc usage (same value on all tasks, or -1) + mem_fresh: bool = False # did this snapshot's mem numbers come from a fresh read? + # cgroup-based metrics (None when not in a systemd slice / no per-task scope) + cgroup_cpu_pct: Optional[float] = None # aggregate CPU % from cgroup cpu.stat + cgroup_mem_mb: Optional[float] = None # aggregate memory from cgroup memory.current + + +def _get_child_procs_fallback(base_pid: int) -> List[int]: + """Pure-bash fallback when psutil.Process.children raises AccessDenied.""" + script = r''' + childprocs() { + local parent=$1 + if [ ! "$2" ]; then child_pid_list=""; fi + if [ "$parent" ]; then + child_pid_list="$child_pid_list $parent" + for childpid in $(pgrep -P ${parent}); do + childprocs $childpid "nottoplevel" + done + fi + if [ ! "$2" ]; then echo "${child_pid_list}"; fi + } + ''' + full = script + f"\nchildprocs {base_pid}\n" + out = subprocess.check_output(full, shell=True) + pids: List[int] = [] + for tok in out.decode().split(): + try: + pids.append(int(tok)) + except ValueError: + continue + return pids + + +class PsutilBackend: + """psutil-based monitor backend. Caches per-pid Process objects so that + cpu_percent(interval=None) has a stable 'previous reading' baseline.""" + + def __init__(self): + if not HAVE_PSUTIL: + raise RuntimeError("psutil not available") + self._proc_cache: Dict[int, "psutil.Process"] = {} + # Prime baseline on the manager process so first delta is sensible. + try: + psutil.cpu_percent(interval=None) + except Exception: + pass + + def _get_or_add(self, pid: int) -> Optional["psutil.Process"]: + p = self._proc_cache.get(pid) + if p is not None: + return p + try: + p = psutil.Process(pid) + # Deliberately DO NOT call p.cpu_percent(interval=None) here to prime. + # psutil's rule: the first call after Process() construction returns + # a CPU % relative to process creation -- useful as a "since start" + # figure, not "since last poll". On multi-process tasks that spawn + # children between monitor ticks, priming right before the read + # yields 0 because no time has elapsed. Letting the monitor thread's + # regular cadence handle both calls produces correct deltas from + # tick 2 onward (tick 1 returns the since-creation value, which is + # a reasonable initial estimate anyway). + self._proc_cache[pid] = p + return p + except Exception: + return None + + def forget(self, pid: int) -> None: + self._proc_cache.pop(pid, None) + + def sweep_dead(self) -> int: + """Evict cached Process objects whose underlying PID is gone. + + Called once per monitor pass. Keeps the cache size bounded over + long runs where short-lived DPL/FAIRMQ children come and go. + """ + dead = 0 + for pid in list(self._proc_cache.keys()): + try: + if not self._proc_cache[pid].is_running(): + del self._proc_cache[pid] + dead += 1 + except Exception: + self._proc_cache.pop(pid, None) + dead += 1 + return dead + + def sample( + self, + root_pid: int, + want_mem: bool, + ) -> Tuple[float, float, float, float, int]: + """Return (cpu_pct_sum, pss_mb, uss_mb, swap_mb, nice_of_root). + + Sums over root and all descendants. If want_mem is False the + memory figures are returned as 0 (caller should interpret -> use + last known). + """ + root = self._get_or_add(root_pid) + if root is None: + return 0.0, 0.0, 0.0, 0.0, 0 + + # Enumerate the PIDs of the whole process tree, then look each one + # up in the cache. CRITICAL: psutil.Process.children() returns NEW + # Process objects on every call, which means their cpu_percent() + # baseline is reset each tick -> we'd read 0.0 forever. Always go + # through _get_or_add so baselines persist between ticks. + try: + child_pids = [c.pid for c in root.children(recursive=True)] + except (psutil.NoSuchProcess,): + return 0.0, 0.0, 0.0, 0.0, 0 + except (psutil.AccessDenied, PermissionError): + try: + child_pids = _get_child_procs_fallback(root_pid) + if root_pid in child_pids: + child_pids.remove(root_pid) + except Exception: + child_pids = [] + + procs = [root] + for pid in child_pids: + p = self._get_or_add(pid) + if p is not None: + procs.append(p) + + cpu_sum = 0.0 + pss_sum = 0.0 + uss_sum = 0.0 + swap_sum = 0.0 + for p in procs: + # CPU: cheap + try: + cpu_sum += p.cpu_percent(interval=None) + except (psutil.NoSuchProcess, psutil.AccessDenied): + pass + + if want_mem: + try: + mi = p.memory_full_info() + pss_sum += getattr(mi, "pss", 0) or 0 + uss_sum += getattr(mi, "uss", 0) or 0 + swap_sum += getattr(mi, "swap", 0) or 0 + except (psutil.NoSuchProcess, psutil.AccessDenied): + pass + + # bytes -> MB + pss_mb = pss_sum / 1024.0 / 1024.0 + uss_mb = uss_sum / 1024.0 / 1024.0 + swap_mb = swap_sum / 1024.0 / 1024.0 + + try: + nice = root.nice() + except Exception: + nice = 0 + + return cpu_sum, pss_mb, uss_mb, swap_mb, nice + + +def _read_cgroup_v2_dir(pid: int) -> Optional[str]: + """Return the cgroup v2 directory for *pid*, or None if not on cgroup v2.""" + try: + with open(f"/proc/{pid}/cgroup") as fh: + for line in fh: + parts = line.strip().split(":", 2) + # cgroup v2 unified hierarchy: single entry "0::" + if len(parts) == 3 and parts[0] == "0": + rel = parts[2].lstrip("/") + candidate = f"/sys/fs/cgroup/{rel}" if rel else "/sys/fs/cgroup" + if os.path.isdir(candidate): + return candidate + except OSError: + pass + return None + + +class CgroupV2Monitor: + """Reads aggregate CPU time and memory for one cgroup v2 directory. + + Works with cgroup v2 (unified hierarchy) only — the format used by + systemd on modern Linux. Instantiating this class is always safe; call + ``available`` to test whether a valid cgroup path was found/given before + using ``sample()``. + + If *cgroup_dir* is provided the directory is used directly (for per-task + monitoring of a known scope). Otherwise the caller's own cgroup is + detected from ``/proc/self/cgroup`` (for the global runner-level monitor). + + CPU is computed by differentiating the ``usage_usec`` counter in + ``cpu.stat``. The first ``sample()`` call primes the counter and returns + ``None`` for cpu_pct; subsequent calls give the average utilisation + (as a percentage of one core, same units as psutil's cpu_percent). + + Memory is read from ``memory.current`` (bytes → MB). It counts all + memory mapped by processes in the cgroup including file-backed pages; + shared pages are counted once per cgroup, not per process. For an + apples-to-apples comparison with psutil's PSS metric use the + ``memory.stat`` anon field (not implemented here — ``memory.current`` + is the right proxy for a hard MemoryMax enforcement budget). + """ + + def __init__(self, cgroup_dir: Optional[str] = None) -> None: + if cgroup_dir is not None: + self._cgroup_dir: Optional[str] = cgroup_dir if os.path.isdir(cgroup_dir) else None + else: + self._cgroup_dir = _read_cgroup_v2_dir(os.getpid()) + self._last_usage_usec: Optional[int] = None + self._last_ts: Optional[float] = None + + @property + def available(self) -> bool: + return self._cgroup_dir is not None + + def sample(self) -> Tuple[Optional[float], Optional[float]]: + """Return (cpu_pct, mem_mb). cpu_pct is None on the first call.""" + if not self._cgroup_dir: + return None, None + + cpu_pct: Optional[float] = None + mem_mb: Optional[float] = None + + # --- CPU --- + try: + with open(os.path.join(self._cgroup_dir, "cpu.stat")) as fh: + for line in fh: + if line.startswith("usage_usec"): + usage_usec = int(line.split()[1]) + now = time.monotonic() + if self._last_usage_usec is not None and self._last_ts is not None: + dt = now - self._last_ts + if dt > 0: + d_usec = usage_usec - self._last_usage_usec + # d_usec / (dt * 1e6) is CPU fraction relative to 1 core; + # multiply by 100 to match psutil's cpu_percent scale. + cpu_pct = (d_usec / (dt * 1e6)) * 100.0 + self._last_usage_usec = usage_usec + self._last_ts = now + break + except OSError: + pass + + # --- Memory --- + try: + with open(os.path.join(self._cgroup_dir, "memory.current")) as fh: + mem_mb = int(fh.read().strip()) / 1024.0 / 1024.0 + except OSError: + pass + + return cpu_pct, mem_mb + + +class MonitorThread(threading.Thread): + """Background thread polling a set of (tid, pid) pairs. + + The scheduler registers tasks via register()/deregister(). On each + pass the thread computes the latest snapshot and stores it in + self.snapshots; the scheduler reads atomically. + + CPU and MEM have independent cadences because PSS reads via + /proc//smaps_rollup are ~10x more expensive than cpu_percent. + """ + + def __init__( + self, + cpu_interval: float = 1.0, + mem_interval: float = 5.0, + backend: Optional[PsutilBackend] = None, + monitor_disc: bool = False, + disc_path: str = ".", + global_cgroup_dir: Optional[str] = None, + ): + super().__init__(daemon=True, name="o2dpg-monitor") + self.cpu_interval = cpu_interval + self.mem_interval = mem_interval + self.backend = backend if backend is not None else PsutilBackend() + self.monitor_disc = monitor_disc + self.disc_path = disc_path + + self._lock = threading.Lock() + # registered[tid] = {"pid": int, "name": str, "labels": list, + # "start_time": float} + self._registered: Dict[int, Dict] = {} + self._snapshots: Dict[int, TaskSnapshot] = {} + self._stop_event = threading.Event() + self._last_mem_ts: float = 0.0 + self._last_disc_mb: float = -1.0 + self._last_disc_ts: float = 0.0 + # Tick counter: incremented once per monitor pass (roughly once per + # cpu_interval). Read by the executor when writing metric lines so + # the "iter" field reflects wall-clock ticks like the prototype did. + self.tick: int = 0 + + # Opportunistic cgroup v2 global monitor. Only enabled when the + # caller passes an explicit *slice* directory (i.e. the runner was + # launched under --systemd-run). Without an explicit directory we + # skip cgroup monitoring entirely — the runner's own cgroup outside a + # dedicated slice covers unrelated user-session processes and gives + # misleading aggregates. + # Written only from the monitor thread; read from the executor thread — + # GIL makes bare float/None assignment atomic. + self._cgroup: Optional[CgroupV2Monitor] = ( + CgroupV2Monitor(cgroup_dir=global_cgroup_dir) if global_cgroup_dir else None + ) + self.global_cpu_pct: Optional[float] = None # cgroup-aggregate CPU % + self.global_mem_mb: Optional[float] = None # cgroup-aggregate memory MB + if self._cgroup is not None and self._cgroup.available: + log.info("CgroupV2Monitor global active at %s", self._cgroup._cgroup_dir) + + # ----- registration ----- + def register( + self, + tid: int, + pid: int, + name: str, + labels: List[str], + start_time: float, + resolve_cgroup: bool = False, + ) -> None: + """Register a task for monitoring. + + If *resolve_cgroup* is True the monitor thread will lazily locate the + task's cgroup v2 directory by inspecting the first child process of + *pid* (which is the systemd-run wrapper when per-task scopes are used). + Once resolved a per-task CgroupV2Monitor is created and its readings + are stored in the TaskSnapshot alongside the psutil figures. + """ + with self._lock: + self._registered[tid] = { + "pid": pid, + "name": name, + "labels": labels, + "start_time": start_time, + "resolve_cgroup": resolve_cgroup, + "cgroup_monitor": None, # filled lazily by _one_pass + } + + def deregister(self, tid: int) -> None: + with self._lock: + entry = self._registered.pop(tid, None) + if entry is not None: + self.backend.forget(entry["pid"]) + + # ----- snapshot access ----- + def latest(self) -> Dict[int, TaskSnapshot]: + """Return a shallow copy of current snapshots.""" + with self._lock: + return dict(self._snapshots) + + def latest_for(self, tid: int) -> Optional[TaskSnapshot]: + with self._lock: + return self._snapshots.get(tid) + + # ----- control ----- + def stop(self) -> None: + self._stop_event.set() + + # ----- loop ----- + def _disc_usage_mb(self) -> float: + try: + out = subprocess.check_output(["du", "-sb", self.disc_path], text=True) + return int(out.split()[0]) / 1024.0 / 1024.0 + except Exception: + return -1.0 + + def _one_pass(self, now: float) -> None: + self.tick += 1 + + # cgroup global totals (cheap reads, always done every pass) + if self._cgroup is not None and self._cgroup.available: + g_cpu, g_mem = self._cgroup.sample() + if g_cpu is not None: + self.global_cpu_pct = g_cpu + if g_mem is not None: + self.global_mem_mb = g_mem + + want_mem = (now - self._last_mem_ts) >= self.mem_interval + if want_mem: + self._last_mem_ts = now + + if self.monitor_disc and (now - self._last_disc_ts) >= self.mem_interval: + self._last_disc_mb = self._disc_usage_mb() + self._last_disc_ts = now + + # snapshot registrations under lock, release for the slow work + with self._lock: + registered_copy = dict(self._registered) + + new_snaps: Dict[int, TaskSnapshot] = {} + for tid, info in registered_copy.items(): + pid = info["pid"] + cpu_pct, pss_mb, uss_mb, swap_mb, nice = self.backend.sample(pid, want_mem) + # fall back to previous mem reading if not a mem-interval tick + prev = self._snapshots.get(tid) + if not want_mem and prev is not None: + pss_mb = prev.pss_mb + uss_mb = prev.uss_mb + swap_mb = prev.swap_mb + + # --- per-task cgroup monitoring --- + # Lazy resolution: when the task runs inside a systemd scope, its + # direct child of the systemd-run wrapper PID lands in that scope's + # cgroup. We probe once per pass until the child appears. + if info.get("resolve_cgroup") and info.get("cgroup_monitor") is None: + try: + p_obj = self.backend._get_or_add(pid) + if p_obj is not None: + children = p_obj.children() + if children: + cgroup_dir = _read_cgroup_v2_dir(children[0].pid) + if cgroup_dir: + info["cgroup_monitor"] = CgroupV2Monitor(cgroup_dir=cgroup_dir) + log.info("Per-task cgroup resolved: tid=%d %s → %s", + tid, info["name"], cgroup_dir) + except Exception: + pass + + cgroup_cpu: Optional[float] = None + cgroup_mem: Optional[float] = None + cm: Optional[CgroupV2Monitor] = info.get("cgroup_monitor") + if cm is not None and cm.available: + cgroup_cpu, cgroup_mem = cm.sample() + # cgroup_cpu is None on the very first sample() call (no prior + # baseline yet) and whenever the scope's cpu.stat is unreadable + # (e.g. after --collect removes the finished scope). In both + # cases we leave it as None rather than carrying forward a stale + # value — one tick with None is preferable to a wrong number. + # ------------------------------------------------------- + + t_delta_ms = int((now - info["start_time"]) * 1000) + new_snaps[tid] = TaskSnapshot( + tid=tid, name=info["name"], + t_delta_ms=t_delta_ms, cpu_pct=cpu_pct, + uss_mb=uss_mb, pss_mb=pss_mb, swap_mb=swap_mb, + nice=nice, labels=info["labels"], + disc_mb=self._last_disc_mb, + mem_fresh=want_mem, + cgroup_cpu_pct=cgroup_cpu, + cgroup_mem_mb=cgroup_mem, + ) + + with self._lock: + self._snapshots = new_snaps + + # housekeeping: evict dead cached Process objects once per pass + try: + self.backend.sweep_dead() + except Exception: + pass + + def run(self) -> None: + log.debug("Monitor thread started (cpu=%.2fs, mem=%.2fs)", + self.cpu_interval, self.mem_interval) + while not self._stop_event.is_set(): + now = time.perf_counter() + try: + self._one_pass(now) + except Exception: + log.exception("Monitor pass failed (continuing)") + # Sleep at the CPU cadence; MEM is gated by its own interval. + self._stop_event.wait(self.cpu_interval) + log.debug("Monitor thread exiting") diff --git a/MC/workflow_runner/o2dpg_runner/resources.py b/MC/workflow_runner/o2dpg_runner/resources.py new file mode 100644 index 000000000..4b187d90c --- /dev/null +++ b/MC/workflow_runner/o2dpg_runner/resources.py @@ -0,0 +1,362 @@ +"""Resource management: per-task estimates, global budget, semaphores. + +Structurally equivalent to the prototype's TaskResources / ResourceManager, +with three changes: + 1. is_within_limits() actually checks MEM against mem_limit (not CPU). + 2. No module-level args dependency; n_backfill is passed in. + 3. Booking/unbooking is explicit about which bucket (default / backfill). + +Sibling-sampling for --dynamic-resources still runs inside unbook(), same +ordering as before. +""" + +from __future__ import annotations + +import logging +import os +from dataclasses import dataclass, field +from typing import Dict, Iterator, List, Optional, Tuple + +log = logging.getLogger(__name__) + + +class Semaphore: + """A named mutual-exclusion flag shared by a group of tasks. + + Deliberately not threading.Semaphore -- this is a logical gate checked + by the scheduler, not a concurrency primitive. + """ + __slots__ = ("locked",) + + def __init__(self): + self.locked = False + + def lock(self): + self.locked = True + + def unlock(self): + self.locked = False + + +@dataclass +class ResourceBoundaries: + cpu_limit: float + mem_limit: float + dynamic_resources: bool = False + optimistic_resources: bool = False + + +class TaskResources: + """Resource accounting for a single task.""" + + def __init__( + self, + tid: int, + name: str, + cpu: float, + cpu_relative: Optional[float], + mem: float, + boundaries: ResourceBoundaries, + ): + self.tid = tid + self.name = name + # originals (never mutated) + self.cpu_assigned_original = cpu + self.mem_assigned_original = mem + self.cpu_relative = cpu_relative if cpu_relative else 1.0 + # transient assignments (may be updated by sampling) + self.cpu_assigned = cpu + self.mem_assigned = mem + self.boundaries = boundaries + # sampled (after a sibling finished) + self.cpu_sampled: Optional[float] = None + self.mem_sampled: Optional[float] = None + # live monitor feed + self.time_collect: List[float] = [] + self.cpu_collect: List[float] = [] + self.mem_collect: List[float] = [] + # siblings (same "global" task name) + self.related_tasks: Optional[List["TaskResources"]] = None + self.semaphore: Optional[Semaphore] = None + self.nice_value: Optional[int] = None + self.booked = False + + # ----- helpers ----- + @property + def is_done(self) -> bool: + return bool(self.time_collect) and not self.booked + + def is_within_limits(self) -> bool: + """Check the current assignment against global boundaries.""" + ok_cpu = self.cpu_assigned <= self.boundaries.cpu_limit + ok_mem = self.mem_assigned <= self.boundaries.mem_limit + if not ok_cpu: + log.warning("CPU of %s exceeds limit: %.2f > %.2f", + self.name, self.cpu_assigned, self.boundaries.cpu_limit) + if not ok_mem: + log.warning("MEM of %s exceeds limit: %.2f > %.2f", + self.name, self.mem_assigned, self.boundaries.mem_limit) + return ok_cpu and ok_mem + + def limit_resources(self, cpu_limit: float = None, mem_limit: float = None) -> None: + if cpu_limit is None: + cpu_limit = self.boundaries.cpu_limit + if mem_limit is None: + mem_limit = self.boundaries.mem_limit + self.cpu_assigned = min(self.cpu_assigned, cpu_limit) + self.mem_assigned = min(self.mem_assigned, mem_limit) + + def add_sample(self, time_passed: float, cpu_fraction: float, mem_mb: float) -> None: + """Record a monitor sample.""" + self.time_collect.append(time_passed) + self.cpu_collect.append(cpu_fraction) + self.mem_collect.append(mem_mb) + + def sample_resources(self) -> None: + """Compute CPU/MEM sample and propagate to un-started siblings.""" + if not self.is_done: + return + + if len(self.time_collect) < 3: + self.cpu_sampled = self.cpu_assigned + self.mem_sampled = self.mem_assigned + log.debug("Task %s: not enough samples (<3); using assigned as sampled", + self.name) + else: + # Weighted mean of CPU over sample intervals, skipping the first + # (cpu_percent(interval=None) first reading is meaningless). + deltas = [self.time_collect[i + 1] - self.time_collect[i] + for i in range(len(self.time_collect) - 1)] + tot = sum(deltas) or 1.0 + cpu_integral = sum( + c * dt for c, dt in zip(self.cpu_collect[1:], deltas) if c >= 0 + ) + self.cpu_sampled = cpu_integral / tot + self.mem_sampled = max(self.mem_collect) + + if self.related_tasks is None: + return + + # aggregate over finished siblings + mem_agg = 0.0 + cpu_list: List[float] = [] + for sib in self.related_tasks: + if sib.is_done and sib.cpu_sampled is not None and sib.mem_sampled is not None: + mem_agg = max(mem_agg, sib.mem_sampled) + cpu_list.append(sib.cpu_sampled) + if not cpu_list: + return + cpu_agg = sum(cpu_list) / len(cpu_list) + + if cpu_agg > self.boundaries.cpu_limit: + log.warning("Sampled CPU (%.2f) exceeds limit (%.2f)", + cpu_agg, self.boundaries.cpu_limit) + elif cpu_agg <= 0: + # a zero reading is missing information, not a task that needs no + # CPU; handing it on would let every sibling be admitted at once + log.debug("Sampled CPU<=0 for %s; reverting to assigned", self.name) + cpu_agg = self.cpu_assigned + + if mem_agg > self.boundaries.mem_limit: + log.warning("Sampled MEM (%.2f) exceeds limit (%.2f)", + mem_agg, self.boundaries.mem_limit) + elif mem_agg <= 0: + log.debug("Sampled MEM<=0 for %s; reverting to assigned", self.name) + mem_agg = self.mem_assigned + + for sib in self.related_tasks: + if sib.is_done or sib.booked: + continue + sib.cpu_assigned = cpu_agg * sib.cpu_relative + sib.mem_assigned = mem_agg + sib.limit_resources() + + +class ResourceManager: + """Central accounting: who is booked, which bucket, what's left.""" + + def __init__( + self, + cpu_limit: float, + mem_limit: float, + procs_parallel_max: int = 100, + n_backfill_max: int = 1, + backfill_cpu_factor: float = 1.5, + backfill_mem_factor: float = 1.5, + dynamic_resources: bool = False, + optimistic_resources: bool = False, + ): + self.boundaries = ResourceBoundaries( + cpu_limit, mem_limit, dynamic_resources, optimistic_resources + ) + self.resources: List[TaskResources] = [] + self._related_by_name: Dict[str, List[TaskResources]] = {} + self._semaphores: Dict[str, Semaphore] = {} + + # default-priority bucket + self.cpu_booked = 0.0 + self.mem_booked = 0.0 + self.n_procs = 0 + + # backfill (niced) bucket + self.cpu_booked_backfill = 0.0 + self.mem_booked_backfill = 0.0 + self.n_procs_backfill = 0 + + self.procs_parallel_max = procs_parallel_max + self.n_backfill_max = n_backfill_max + self.backfill_cpu_factor = backfill_cpu_factor + self.backfill_mem_factor = backfill_mem_factor + + try: + self.nice_default = os.nice(0) + except (AttributeError, OSError): + self.nice_default = 0 + self.nice_backfill = self.nice_default + 19 + + # ----- registration ----- + def add_task( + self, + name: str, + related_name: Optional[str], + cpu: float, + cpu_relative: Optional[float], + mem: float, + semaphore_string: Optional[str] = None, + ) -> TaskResources: + res = TaskResources( + len(self.resources), name, cpu, cpu_relative, mem, self.boundaries + ) + if not res.is_within_limits() and not self.boundaries.optimistic_resources: + raise ResourceLimitExceeded( + f"Task {name} exceeds resource boundaries " + f"(cpu={cpu}/{self.boundaries.cpu_limit}, " + f"mem={mem}/{self.boundaries.mem_limit}). " + f"Use --optimistic-resources to attempt anyway." + ) + res.limit_resources() + self.resources.append(res) + + if semaphore_string: + if semaphore_string not in self._semaphores: + self._semaphores[semaphore_string] = Semaphore() + res.semaphore = self._semaphores[semaphore_string] + + if related_name: + bucket = self._related_by_name.setdefault(related_name, []) + bucket.append(res) + res.related_tasks = bucket + + return res + + # ----- monitor hook ----- + def add_monitored(self, tid: int, t_delta: float, cpu_fraction: float, mem_mb: float) -> None: + self.resources[tid].add_sample(t_delta, cpu_fraction, mem_mb) + + # ----- booking ----- + def book(self, tid: int, nice_value: int) -> None: + res = self.resources[tid] + # Prior check is expected to have set nice_value; if not, force backfill. + if res.nice_value is None: + log.warning("Task %d booked without prior ok_to_submit check; forcing backfill", + tid) + nice_value = self.nice_backfill + + res.nice_value = nice_value + res.booked = True + if res.semaphore is not None: + res.semaphore.lock() + if nice_value != self.nice_default: + self.n_procs_backfill += 1 + self.cpu_booked_backfill += res.cpu_assigned + self.mem_booked_backfill += res.mem_assigned + else: + self.n_procs += 1 + self.cpu_booked += res.cpu_assigned + self.mem_booked += res.mem_assigned + + def unbook(self, tid: int) -> None: + res = self.resources[tid] + res.booked = False + if self.boundaries.dynamic_resources: + res.sample_resources() + if res.semaphore is not None: + res.semaphore.unlock() + if res.nice_value != self.nice_default: + self.cpu_booked_backfill -= res.cpu_assigned + self.mem_booked_backfill -= res.mem_assigned + self.n_procs_backfill -= 1 + if self.n_procs_backfill <= 0: + self.cpu_booked_backfill = 0.0 + self.mem_booked_backfill = 0.0 + else: + self.n_procs -= 1 + self.cpu_booked -= res.cpu_assigned + self.mem_booked -= res.mem_assigned + if self.n_procs <= 0: + self.cpu_booked = 0.0 + self.mem_booked = 0.0 + + # ----- queries ----- + def total_procs(self) -> int: + return self.n_procs + self.n_procs_backfill + + def at_proc_cap(self) -> bool: + return self.total_procs() >= self.procs_parallel_max + + def cpu_free_default(self) -> float: + return self.boundaries.cpu_limit - self.cpu_booked + + def mem_free_default(self) -> float: + return self.boundaries.mem_limit - self.mem_booked + + def fits_default(self, res: TaskResources) -> bool: + return ( + self.cpu_booked + res.cpu_assigned <= self.boundaries.cpu_limit + and self.mem_booked + res.mem_assigned <= self.boundaries.mem_limit + ) + + def fits_backfill( + self, + res: TaskResources, + cpu_factor: float = None, + mem_factor: float = None, + ) -> bool: + if cpu_factor is None: + cpu_factor = self.backfill_cpu_factor + if mem_factor is None: + mem_factor = self.backfill_mem_factor + if self.n_procs_backfill >= self.n_backfill_max: + return False + # don't backfill with huge tasks (originals: avoid tasks too close to limit) + if res.cpu_assigned > 0.9 * self.boundaries.cpu_limit: + return False + # mem per core sanity: don't launch something whose mem is huge relative to + # the CPU budget (original heuristic: mem/cpu_limit >= 1900). + if self.boundaries.cpu_limit > 0 and res.mem_assigned / self.boundaries.cpu_limit >= 1900: + return False + + ok_cpu = (self.cpu_booked_backfill + res.cpu_assigned + <= self.boundaries.cpu_limit) + ok_cpu = ok_cpu and ( + self.cpu_booked + self.cpu_booked_backfill + res.cpu_assigned + <= cpu_factor * self.boundaries.cpu_limit + ) + ok_mem = ( + self.mem_booked + self.mem_booked_backfill + res.mem_assigned + <= mem_factor * self.boundaries.mem_limit + ) + return ok_cpu and ok_mem + + def can_be_submitted_at_all(self, res: TaskResources) -> bool: + """True if a task is not blocked by its semaphore and is not already booked.""" + if res.booked: + return False + if res.semaphore is not None and res.semaphore.locked: + return False + return True + + +class ResourceLimitExceeded(Exception): + """Raised when a task's declared resources exceed the global boundaries + and --optimistic-resources was not given.""" diff --git a/MC/workflow_runner/o2dpg_runner/scheduler/__init__.py b/MC/workflow_runner/o2dpg_runner/scheduler/__init__.py new file mode 100644 index 000000000..0386a93ea --- /dev/null +++ b/MC/workflow_runner/o2dpg_runner/scheduler/__init__.py @@ -0,0 +1,22 @@ +from .base import SchedulerPolicy, SchedulerState +from .timeframe import TimeframeFirstPolicy +from .critical_path import CriticalPathPolicy +from .best_fit import BestFitBackfillPolicy + + +def get_policy(name: str) -> SchedulerPolicy: + name = name.lower().strip() + if name in ("timeframe", "tf", "legacy"): + return TimeframeFirstPolicy() + if name in ("critical-path", "cp"): + return CriticalPathPolicy() + if name in ("best-fit", "bf"): + return BestFitBackfillPolicy() + raise ValueError(f"unknown scheduler policy: {name}") + + +__all__ = [ + "SchedulerPolicy", "SchedulerState", + "TimeframeFirstPolicy", "CriticalPathPolicy", "BestFitBackfillPolicy", + "get_policy", +] diff --git a/MC/workflow_runner/o2dpg_runner/scheduler/base.py b/MC/workflow_runner/o2dpg_runner/scheduler/base.py new file mode 100644 index 000000000..9ea3fea36 --- /dev/null +++ b/MC/workflow_runner/o2dpg_runner/scheduler/base.py @@ -0,0 +1,62 @@ +"""Pluggable scheduler policy interface. + +A SchedulerPolicy has two responsibilities: + 1. order(candidates, state) -> ordered list of tids + (how to prioritize among runnable tasks) + 2. pick_submittable(ordered, resource_manager) -> iterator of (tid, nice) + (which of the ordered tasks actually fit in the remaining budget, + at what nice level, subject to policy-specific packing rules) + +Resources-aware bookkeeping (what's currently booked, what the limits +are) lives in ResourceManager. Policies only read it; they don't mutate +it. The executor calls rm.book() after picking a task. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Iterator, List, Tuple, TYPE_CHECKING + +if TYPE_CHECKING: + from ..resources import ResourceManager + + +@dataclass +class SchedulerState: + """Everything a policy might want to know about the current run. + + Kept lightweight: policies that don't need a field just ignore it. + """ + # Static data + timeframe_of: List[int] = field(default_factory=list) # tid -> timeframe + descendants_count: List[int] = field(default_factory=list) # |desc(tid)| + critical_path: List[float] = field(default_factory=list) # longest_path_length weighted by walltime (or cpu fallback) + task_cpu: List[float] = field(default_factory=list) + task_mem: List[float] = field(default_factory=list) + task_walltime: List[float] = field(default_factory=list) # per-task walltime [s]; 0 if unknown + # Derived weight tuples (static, cached by executor) + # timeframe_weight[tid] = (timeframe, -num_descendants) + timeframe_weight: List[Tuple[int, int]] = field(default_factory=list) + + +class SchedulerPolicy: + """Base class. Subclasses override order() and/or pick_submittable().""" + + name = "base" + + def order(self, candidates: List[int], state: SchedulerState) -> List[int]: + """Return candidates ordered by policy preference (best first).""" + raise NotImplementedError + + def pick_submittable( + self, + ordered: List[int], + rm: "ResourceManager", + ) -> Iterator[Tuple[int, int]]: + """Yield (tid, nice_value) for tasks that fit now. + + Removes nothing; the executor is responsible for calling rm.book() + and rebuilding candidate lists. pick_submittable is a pure + read-only view onto rm's current state. + """ + raise NotImplementedError diff --git a/MC/workflow_runner/o2dpg_runner/scheduler/best_fit.py b/MC/workflow_runner/o2dpg_runner/scheduler/best_fit.py new file mode 100644 index 000000000..a6546ac9d --- /dev/null +++ b/MC/workflow_runner/o2dpg_runner/scheduler/best_fit.py @@ -0,0 +1,104 @@ +"""Best-fit bin-packing scheduler. + +Instead of following the ordered list linearly, it picks the candidate +that maximizes a fitness score given the current remaining budget. +Iterates until nothing fits. The backfill pass uses the same idea. + +Fitness = critical_path_weight * packing_tightness, where: + - critical_path_weight is state.critical_path[tid] — remaining walltime + on the longest path (walltime-weighted when learned data is available, + cpu-weighted otherwise). This is the "hybrid CP + packing" heuristic: + prefer tasks on the longest critical path, but among those that fit + similarly, pick the one that uses capacity most tightly. + - packing_tightness = max(c/cpu_free, m/mem_free) — dominant-resource + utilisation fraction. A task scores high when it fills at least one + resource bin well; unlike 1/max(ratio), this is not dominated by a + resource with extreme slack (e.g. 60 GB mem_limit with tasks using 1 GB). + +Default ordering is still critical-path (good baseline); pick_submittable +re-ranks on the fly within the fitting set. +""" + +from __future__ import annotations + +from typing import Iterator, List, Optional, Tuple + +from .base import SchedulerPolicy, SchedulerState +from .critical_path import CriticalPathPolicy +from ..resources import ResourceManager + + +class BestFitBackfillPolicy(SchedulerPolicy): + name = "best-fit" + + def __init__(self): + self._ordering = CriticalPathPolicy() + self._state: Optional[SchedulerState] = None + + def order(self, candidates: List[int], state: SchedulerState) -> List[int]: + self._state = state + return self._ordering.order(candidates, state) + + @staticmethod + def _fitness(res, state: SchedulerState, cpu_free: float, mem_free: float) -> float: + """Higher is better; negative if it doesn't fit.""" + c = max(res.cpu_assigned, 0.01) + m = max(res.mem_assigned, 1.0) + cpu_ratio = cpu_free / c + mem_ratio = mem_free / m + if cpu_ratio < 1 or mem_ratio < 1: + return -1.0 + # Dominant-resource utilisation: fraction of the most-used resource. + # max(c/cpu_free, m/mem_free) = max(1/cpu_ratio, 1/mem_ratio). + # Higher means the task fills at least one resource bin well. + # This is correct in both constrained and resource-ample (serial) modes: + # unlike 1/max(ratio), it is not dominated by a resource with extreme slack. + tightness = max(c / cpu_free, m / mem_free) + # CP weight: remaining walltime on the longest path from this task. + # Rewards placing tasks that unblock the most remaining work first. + # Falls back to descendants_count+1 when no critical_path available. + cp = state.critical_path[res.tid] if state.critical_path else 0.0 + cp_weight = cp if cp > 0 else (state.descendants_count[res.tid] + 1 if state.descendants_count else 1) + return cp_weight * tightness + + def pick_submittable( + self, ordered: List[int], rm: ResourceManager + ) -> Iterator[Tuple[int, int]]: + if rm.at_proc_cap(): + return + + state = self._state + assert state is not None, "order() must be called before pick_submittable()" + + # --- default-nice: iterative best-fit --- + pool = list(ordered) + while pool and not rm.at_proc_cap(): + cpu_free = rm.cpu_free_default() + mem_free = rm.mem_free_default() + best_tid = -1 + best_score = -1.0 + for tid in pool: + res = rm.resources[tid] + if not rm.can_be_submitted_at_all(res): + continue + s = self._fitness(res, state, cpu_free, mem_free) + if s > best_score: + best_score = s + best_tid = tid + if best_tid < 0 or best_score < 0: + break + pool.remove(best_tid) + res = rm.resources[best_tid] + res.nice_value = rm.nice_default + yield best_tid, rm.nice_default + + # --- backfill pass --- + if rm.at_proc_cap(): + return + for tid in pool: + res = rm.resources[tid] + if not rm.can_be_submitted_at_all(res): + continue + if rm.fits_backfill(res): + res.nice_value = rm.nice_backfill + yield tid, rm.nice_backfill diff --git a/MC/workflow_runner/o2dpg_runner/scheduler/critical_path.py b/MC/workflow_runner/o2dpg_runner/scheduler/critical_path.py new file mode 100644 index 000000000..d07a6b748 --- /dev/null +++ b/MC/workflow_runner/o2dpg_runner/scheduler/critical_path.py @@ -0,0 +1,60 @@ +"""Critical-path-first scheduler. + +Sorts candidates by state.critical_path[tid] — the longest remaining +path weight to any leaf. The weight is walltime [s] when --update-resources +has been used to inject learned lifetime data (resources.walltime per task); +it falls back to cpu cores otherwise. Both give a valid makespan proxy; +the walltime variant is strictly more accurate for multithreaded tasks. + +Submit discipline: scan ordered list, submit everything that fits at default +nice, then do a second backfill pass at elevated nice. No should_break — +there is no reason to stop scanning once we have committed to CP ordering. +""" + +from __future__ import annotations + +from typing import Iterator, List, Tuple + +from .base import SchedulerPolicy, SchedulerState +from ..resources import ResourceManager + + +class CriticalPathPolicy(SchedulerPolicy): + name = "critical-path" + + def order(self, candidates: List[int], state: SchedulerState) -> List[int]: + cp = state.critical_path + tfw = state.timeframe_weight + # primary: longest path (largest first); tie-break: timeframe, tid + return sorted( + candidates, + key=lambda t: (-cp[t] if cp else 0, tfw[t][0], t), + ) + + def pick_submittable( + self, ordered: List[int], rm: ResourceManager + ) -> Iterator[Tuple[int, int]]: + if rm.at_proc_cap(): + return + skipped: List[int] = [] + for tid in ordered: + if rm.at_proc_cap(): + return + res = rm.resources[tid] + if not rm.can_be_submitted_at_all(res): + continue + if rm.fits_default(res): + res.nice_value = rm.nice_default + yield tid, rm.nice_default + else: + skipped.append(tid) + + if rm.at_proc_cap(): + return + for tid in skipped: + res = rm.resources[tid] + if not rm.can_be_submitted_at_all(res): + continue + if rm.fits_backfill(res): + res.nice_value = rm.nice_backfill + yield tid, rm.nice_backfill diff --git a/MC/workflow_runner/o2dpg_runner/scheduler/timeframe.py b/MC/workflow_runner/o2dpg_runner/scheduler/timeframe.py new file mode 100644 index 000000000..22df6acd2 --- /dev/null +++ b/MC/workflow_runner/o2dpg_runner/scheduler/timeframe.py @@ -0,0 +1,81 @@ +"""Timeframe-first scheduler: exact behavior of the prototype. + +Order: + (timeframe, -num_descendants) # small TF first, then most-connected tasks + +Submit: + - First pass: default nice. Scan in order; on first non-fitting task, BREAK + (this is the legacy behavior; it can block light tasks behind a heavy one). + - Second pass: backfill nice. Scan remaining in order; no break on miss. + +This is the default policy so a bare-minimum invocation reproduces the +prototype's scheduling decisions. +""" + +from __future__ import annotations + +import logging +from typing import Iterator, List, Tuple + +from .base import SchedulerPolicy, SchedulerState +from ..resources import ResourceManager + +log = logging.getLogger(__name__) + + +class TimeframeFirstPolicy(SchedulerPolicy): + name = "timeframe" + + def __init__(self, drop_should_break: bool = False): + # When True, the default pass does not break on a non-fitting task + # but keeps scanning, so a light task can slip past a heavy one. + self.drop_should_break = drop_should_break + + def order(self, candidates: List[int], state: SchedulerState) -> List[int]: + # sort prefers small timeframe, then more descendants + return sorted( + candidates, + key=lambda t: (state.timeframe_weight[t][0], -state.timeframe_weight[t][1]), + ) + + def pick_submittable( + self, ordered: List[int], rm: ResourceManager + ) -> Iterator[Tuple[int, int]]: + if rm.at_proc_cap(): + return + + # --- default-nice pass --- + skipped_for_backfill: List[int] = [] + for tid in ordered: + if rm.at_proc_cap(): + return + res = rm.resources[tid] + if not rm.can_be_submitted_at_all(res): + continue + if rm.fits_default(res): + res.nice_value = rm.nice_default + yield tid, rm.nice_default + else: + if self.drop_should_break: + skipped_for_backfill.append(tid) + continue + # legacy behavior: the first non-fit breaks the default pass + skipped_for_backfill.extend( + ordered[ordered.index(tid):] + ) + break + + # --- backfill pass --- + if rm.at_proc_cap(): + return + seen = set() + for tid in skipped_for_backfill: + if tid in seen: + continue + seen.add(tid) + res = rm.resources[tid] + if not rm.can_be_submitted_at_all(res): + continue + if rm.fits_backfill(res): + res.nice_value = rm.nice_backfill + yield tid, rm.nice_backfill diff --git a/MC/workflow_runner/o2dpg_runner/tests/__init__.py b/MC/workflow_runner/o2dpg_runner/tests/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/MC/workflow_runner/o2dpg_runner/tests/fixtures/tiny_workflow.json b/MC/workflow_runner/o2dpg_runner/tests/fixtures/tiny_workflow.json new file mode 100644 index 000000000..251cd5e04 --- /dev/null +++ b/MC/workflow_runner/o2dpg_runner/tests/fixtures/tiny_workflow.json @@ -0,0 +1,104 @@ +{ + "stages": [ + { + "name": "__global_init_task__", + "cmd": "NO-COMMAND", + "needs": [], + "cwd": "./", + "timeframe": -1, + "labels": [], + "resources": {"cpu": 1, "mem": 100, "relative_cpu": null}, + "env": {"FOO": "bar"} + }, + { + "name": "bkg", + "cmd": "echo bkg > out.dat", + "needs": [], + "cwd": "./", + "timeframe": -1, + "labels": ["SIM"], + "resources": {"cpu": 2, "mem": 500, "relative_cpu": null} + }, + { + "name": "sgnsim_1", + "cmd": "echo sgnsim_1", + "needs": ["bkg"], + "cwd": "./tf1", + "timeframe": 1, + "labels": ["SIM"], + "resources": {"cpu": 4, "mem": 2000, "relative_cpu": 1.0} + }, + { + "name": "digi_1", + "cmd": "echo digi_1", + "needs": ["sgnsim_1"], + "cwd": "./tf1", + "timeframe": 1, + "labels": ["DIGI"], + "resources": {"cpu": 2, "mem": 1000, "relative_cpu": 1.0} + }, + { + "name": "reco_1", + "cmd": "echo reco_1", + "needs": ["digi_1"], + "cwd": "./tf1", + "timeframe": 1, + "labels": ["RECO"], + "resources": {"cpu": 4, "mem": 1500, "relative_cpu": 1.0} + }, + { + "name": "qc_1", + "cmd": "echo qc_1", + "needs": ["reco_1"], + "cwd": "./tf1", + "timeframe": 1, + "labels": ["QC"], + "resources": {"cpu": 1, "mem": 500, "relative_cpu": 1.0} + }, + { + "name": "sgnsim_2", + "cmd": "echo sgnsim_2", + "needs": ["bkg"], + "cwd": "./tf2", + "timeframe": 2, + "labels": ["SIM"], + "resources": {"cpu": 4, "mem": 2000, "relative_cpu": 1.0} + }, + { + "name": "digi_2", + "cmd": "echo digi_2", + "needs": ["sgnsim_2"], + "cwd": "./tf2", + "timeframe": 2, + "labels": ["DIGI"], + "resources": {"cpu": 2, "mem": 1000, "relative_cpu": 1.0} + }, + { + "name": "reco_2", + "cmd": "echo reco_2", + "needs": ["digi_2"], + "cwd": "./tf2", + "timeframe": 2, + "labels": ["RECO"], + "resources": {"cpu": 4, "mem": 1500, "relative_cpu": 1.0} + }, + { + "name": "qc_2", + "cmd": "echo qc_2", + "needs": ["reco_2"], + "cwd": "./tf2", + "timeframe": 2, + "labels": ["QC"], + "resources": {"cpu": 1, "mem": 500, "relative_cpu": 1.0} + }, + { + "name": "aod", + "cmd": "echo aod", + "needs": ["reco_1", "reco_2"], + "cwd": "./", + "timeframe": -1, + "labels": ["AOD"], + "resources": {"cpu": 2, "mem": 800, "relative_cpu": 1.0} + } + ] +} diff --git a/MC/workflow_runner/o2dpg_runner/tests/test_cache.py b/MC/workflow_runner/o2dpg_runner/tests/test_cache.py new file mode 100644 index 000000000..3dd68cfb7 --- /dev/null +++ b/MC/workflow_runner/o2dpg_runner/tests/test_cache.py @@ -0,0 +1,119 @@ +import json +import os + +import pytest + +from o2dpg_runner.cache import ( + TaskCache, compute_fingerprint, done_path, fingerprint_path, + remove_done_flag, +) + + +def _make_task(cmd="echo hi", needs=None, env=None): + return { + "name": "t", + "cmd": cmd, + "needs": needs or [], + "env": env or {}, + } + + +def _mkdone(tmpdir, logname="t.log"): + logfile = os.path.join(str(tmpdir), logname) + open(done_path(logfile), "w").close() + return logfile + + +def test_fingerprint_stable(): + t = _make_task("cmd1", needs=["a", "b"]) + f1 = compute_fingerprint(t) + f2 = compute_fingerprint(t) + assert f1 == f2 + + +def test_fingerprint_changes_on_cmd(): + f1 = compute_fingerprint(_make_task("cmd1")) + f2 = compute_fingerprint(_make_task("cmd2")) + assert f1["cmd_hash"] != f2["cmd_hash"] + + +def test_fingerprint_needs_order_insensitive(): + f1 = compute_fingerprint(_make_task(needs=["a", "b"])) + f2 = compute_fingerprint(_make_task(needs=["b", "a"])) + assert f1["needs"] == f2["needs"] + + +def test_cache_off_skips_if_done(tmp_path): + logfile = _mkdone(tmp_path) + c = TaskCache("off") + fp = compute_fingerprint(_make_task()) + assert c.is_done(logfile, fp) is True + + +def test_cache_off_does_not_skip_without_done(tmp_path): + logfile = os.path.join(str(tmp_path), "t.log") + c = TaskCache("off") + fp = compute_fingerprint(_make_task()) + assert c.is_done(logfile, fp) is False + + +def test_cache_lenient_keeps_when_no_sidecar(tmp_path): + logfile = _mkdone(tmp_path) + c = TaskCache("lenient") + fp = compute_fingerprint(_make_task()) + assert c.is_done(logfile, fp) is True + + +def test_cache_strict_invalidates_without_sidecar(tmp_path): + logfile = _mkdone(tmp_path) + c = TaskCache("strict") + fp = compute_fingerprint(_make_task()) + assert c.is_done(logfile, fp) is False + # also cleared the _done file + assert not os.path.exists(done_path(logfile)) + + +def test_cache_lenient_invalidates_on_cmd_change(tmp_path): + logfile = _mkdone(tmp_path) + c = TaskCache("lenient") + old_fp = compute_fingerprint(_make_task("old_cmd")) + c.record(logfile, old_fp) + new_fp = compute_fingerprint(_make_task("new_cmd")) + assert c.is_done(logfile, new_fp) is False + assert not os.path.exists(done_path(logfile)) + + +def test_cache_lenient_tolerates_env_change(tmp_path, caplog): + logfile = _mkdone(tmp_path) + c = TaskCache("lenient") + old_fp = compute_fingerprint(_make_task(env={"ALICE_O2_VERSION": "v1"})) + c.record(logfile, old_fp) + new_fp = compute_fingerprint(_make_task(env={"ALICE_O2_VERSION": "v2"})) + assert c.is_done(logfile, new_fp) is True + + +def test_cache_strict_invalidates_on_env_change(tmp_path): + logfile = _mkdone(tmp_path) + c = TaskCache("strict") + old_fp = compute_fingerprint(_make_task(env={"ALICE_O2_VERSION": "v1"})) + c.record(logfile, old_fp) + new_fp = compute_fingerprint(_make_task(env={"ALICE_O2_VERSION": "v2"})) + assert c.is_done(logfile, new_fp) is False + + +def test_cache_record_off_does_nothing(tmp_path): + logfile = _mkdone(tmp_path) + c = TaskCache("off") + fp = compute_fingerprint(_make_task()) + c.record(logfile, fp) + assert not os.path.exists(fingerprint_path(logfile)) + + +def test_remove_done_flag_removes_both(tmp_path): + logfile = _mkdone(tmp_path) + fp_path = fingerprint_path(logfile) + with open(fp_path, "w") as f: + f.write("{}") + remove_done_flag(logfile) + assert not os.path.exists(done_path(logfile)) + assert not os.path.exists(fp_path) diff --git a/MC/workflow_runner/o2dpg_runner/tests/test_executor_e2e.py b/MC/workflow_runner/o2dpg_runner/tests/test_executor_e2e.py new file mode 100644 index 000000000..fc081a685 --- /dev/null +++ b/MC/workflow_runner/o2dpg_runner/tests/test_executor_e2e.py @@ -0,0 +1,216 @@ +"""End-to-end executor tests using the tiny fixture workflow. + +The commands are plain `echo` calls, so no CVMFS / O2 software needed. +We invoke the CLI's main() the same way the shell entry point would. +""" + +import json +import logging +import os +import shutil +import sys +from types import SimpleNamespace + +import pytest + +from o2dpg_runner.config import RunnerConfig +from o2dpg_runner.workflow import build_workflow, load_json +from o2dpg_runner.executor import WorkflowExecutor + +FIXTURE = os.path.join(os.path.dirname(__file__), "fixtures", "tiny_workflow.json") + + +def _make_logger(name, path): + lg = logging.getLogger(name) + lg.handlers.clear() + lg.addHandler(logging.FileHandler(path)) + lg.setLevel(logging.INFO) + lg.propagate = False + return lg + + +def _prep_workflow_in_tmp(tmp_path): + """Copy the fixture into tmp_path and wrap every cmd with a _done marker, + since the stock O2 taskwrapper isn't available in tests.""" + raw = load_json(FIXTURE) + # Patch each command so that it writes the expected _done file after success. + for t in raw["stages"]: + if t["name"] == "__global_init_task__": + continue + cwd = t.get("cwd", "./") + name = t["name"] + # emulate the taskwrapper: run the cmd, then write /.log_done + t["cmd"] = f'({t["cmd"]}) > {name}.log 2>&1 && touch {name}.log_done' + fixture_path = tmp_path / "wf.json" + fixture_path.write_text(json.dumps(raw)) + return str(fixture_path) + + +def _make_executor(tmp_path, cfg_overrides=None): + os.chdir(str(tmp_path)) + wf_path = _prep_workflow_in_tmp(tmp_path) + cfg = RunnerConfig( + workflowfile=wf_path, + cpu_limit=8, + mem_limit=16000, + maxjobs=100, + monitor_interval_cpu=0.2, + monitor_interval_mem=0.5, + ) + if cfg_overrides: + for k, v in cfg_overrides.items(): + setattr(cfg, k, v) + raw = load_json(wf_path) + wf = build_workflow(raw, cfg.target_tasks, cfg.target_labels) + action_logger = _make_logger("a", str(tmp_path / "act.log")) + metric_logger = _make_logger("m", str(tmp_path / "met.log")) + return WorkflowExecutor(cfg, wf, action_logger, metric_logger) + + +def _run(tmp_path, cfg_overrides=None): + exe = _make_executor(tmp_path, cfg_overrides) + return exe.execute(), exe.wf, tmp_path + + +class _FakeProc: + pid = 4242 + + def poll(self): + return None + + def kill(self): + pass + + +class _FakeMonitor: + """One task, one snapshot, and a tick the test advances by hand.""" + + def __init__(self): + self.tick = 1 + self.global_cpu_pct = None + self.global_mem_mb = None + self.snap = SimpleNamespace( + tid=0, name="t", t_delta_ms=1000, cpu_pct=100.0, uss_mb=1.0, + pss_mb=2.0, swap_mb=0.0, nice=0, labels=[], disc_mb=-1, + cgroup_cpu_pct=None, cgroup_mem_mb=None) + + def latest(self): + return {0: self.snap} + + +def test_a_monitor_tick_is_recorded_once_however_often_it_is_polled(tmp_path): + """The wait loop polls faster than the monitor fires. Recording the same + snapshot twice defeats the 'too few samples' guard in sample_resources().""" + exe = _make_executor(tmp_path) + exe.monitor = _FakeMonitor() + exe.process_list = [(0, _FakeProc())] + + for _ in range(5): + exe.wait_for_any([], []) + assert len(exe.rm.resources[0].time_collect) == 1 + + exe.monitor.tick = 2 + exe.monitor.snap.t_delta_ms = 2000 + exe.wait_for_any([], []) + assert exe.rm.resources[0].time_collect == [1000, 2000] + + +def test_executor_runs_all_tasks(tmp_path): + rc, wf, path = _run(tmp_path) + assert rc is False # no errors + # every task should have produced a _done file + for t in wf.stages: + done_file = path / (t.get("cwd", ".") or ".") / f"{t['name']}.log_done" + assert done_file.exists(), f"missing _done for {t['name']}" + + +def test_executor_respects_target_filter(tmp_path): + rc, wf, path = _run(tmp_path, {"target_tasks": ["qc_1"]}) + assert rc is False + # only bkg, sgnsim_1, digi_1, reco_1, qc_1 should have run + expected = {"bkg", "sgnsim_1", "digi_1", "reco_1", "qc_1"} + got = {t["name"] for t in wf.stages} + assert got == expected + + +def test_executor_skips_done_tasks_on_rerun(tmp_path): + # first run + rc1, _, _ = _run(tmp_path) + assert rc1 is False + + # second run: stages unchanged -> every task should be skipped via _done. + os.chdir(str(tmp_path)) + wf_path = str(tmp_path / "wf.json") + cfg = RunnerConfig( + workflowfile=wf_path, + cpu_limit=8, mem_limit=16000, + ) + raw = load_json(wf_path) + wf = build_workflow(raw, cfg.target_tasks, cfg.target_labels) + act = _make_logger("a2", str(tmp_path / "act2.log")) + met = _make_logger("m2", str(tmp_path / "met2.log")) + exe = WorkflowExecutor(cfg, wf, act, met) + # All tasks should appear skippable up front. + for tid in range(wf.n_tasks()): + assert exe.ok_to_skip(tid), f"{wf.id_to_name[tid]} should be skippable" + + +def test_executor_critical_path_policy(tmp_path): + rc, wf, path = _run(tmp_path, {"scheduler_policy": "critical-path"}) + assert rc is False + assert all((path / (t.get("cwd", ".") or ".") / f"{t['name']}.log_done").exists() + for t in wf.stages) + + +def test_executor_best_fit_policy(tmp_path): + rc, wf, path = _run(tmp_path, {"scheduler_policy": "best-fit"}) + assert rc is False + assert all((path / (t.get("cwd", ".") or ".") / f"{t['name']}.log_done").exists() + for t in wf.stages) + + +def test_executor_drop_should_break(tmp_path): + rc, wf, path = _run(tmp_path, {"drop_should_break": True}) + assert rc is False + + +def test_executor_produce_script(tmp_path): + os.chdir(str(tmp_path)) + wf_path = _prep_workflow_in_tmp(tmp_path) + cfg = RunnerConfig( + workflowfile=wf_path, + cpu_limit=8, mem_limit=16000, + produce_script=str(tmp_path / "run.sh"), + ) + raw = load_json(wf_path) + wf = build_workflow(raw, cfg.target_tasks, cfg.target_labels) + act = _make_logger("a3", str(tmp_path / "act3.log")) + met = _make_logger("m3", str(tmp_path / "met3.log")) + exe = WorkflowExecutor(cfg, wf, act, met) + exe.execute() + text = open(str(tmp_path / "run.sh")).read() + assert "bkg" in text + assert "aod" in text + assert "#!/usr/bin/env bash" in text + + +def test_executor_dry_run(tmp_path): + rc, wf, path = _run(tmp_path, {"dry_run": True}) + assert rc is False + # dry run should NOT create _done files + any_done = any((path / (t.get("cwd", ".") or ".") / f"{t['name']}.log_done").exists() + for t in wf.stages) + assert not any_done + + +def test_executor_cache_lenient_end_to_end(tmp_path): + # first run with lenient cache + rc1, _, _ = _run(tmp_path, {"cache_policy": "lenient"}) + assert rc1 is False + # fingerprint sidecars should exist + for t in ["bkg", "qc_1", "aod"]: + # find cwd + cwd = next(s["cwd"] for s in load_json(str(tmp_path / "wf.json"))["stages"] + if s["name"] == t) + assert os.path.exists(str(tmp_path / cwd / f"{t}.log_done")) + assert os.path.exists(str(tmp_path / cwd / f"{t}.log_done.json")) diff --git a/MC/workflow_runner/o2dpg_runner/tests/test_graph.py b/MC/workflow_runner/o2dpg_runner/tests/test_graph.py new file mode 100644 index 000000000..f0e285d51 --- /dev/null +++ b/MC/workflow_runner/o2dpg_runner/tests/test_graph.py @@ -0,0 +1,90 @@ +import pytest + +from o2dpg_runner.graph import ( + build_adjacency, kahn_topological_order, descendants, ancestors, + longest_path_length, invert_adj, root_nodes, +) + + +def test_adjacency_basic(): + # 0 -> 1 -> 2 + # \-> 3 + fwd, rev, ind = build_adjacency(4, [(0, 1), (1, 2), (1, 3)]) + assert fwd == [[1], [2, 3], [], []] + assert rev == [[], [0], [1], [1]] + assert ind == [0, 1, 1, 1] + + +def test_kahn_simple_chain(): + fwd, _, ind = build_adjacency(3, [(0, 1), (1, 2)]) + assert kahn_topological_order(3, fwd, ind) == [0, 1, 2] + + +def test_kahn_diamond(): + # 0 -> 1, 0 -> 2, 1 -> 3, 2 -> 3 + fwd, _, ind = build_adjacency(4, [(0, 1), (0, 2), (1, 3), (2, 3)]) + order = kahn_topological_order(4, fwd, ind) + assert order[0] == 0 + assert order[-1] == 3 + assert set(order[1:3]) == {1, 2} + + +def test_kahn_detects_cycle(): + fwd, _, ind = build_adjacency(3, [(0, 1), (1, 2), (2, 0)]) + with pytest.raises(ValueError): + kahn_topological_order(3, fwd, ind) + + +def test_descendants_simple(): + # 0 -> 1 -> 2 + fwd, _, _ = build_adjacency(3, [(0, 1), (1, 2)]) + assert descendants(fwd, 0) == {1, 2} + assert descendants(fwd, 1) == {2} + assert descendants(fwd, 2) == set() + + +def test_descendants_diamond_memoized(): + fwd, _, _ = build_adjacency(4, [(0, 1), (0, 2), (1, 3), (2, 3)]) + cache = {} + assert descendants(fwd, 0, cache) == {1, 2, 3} + # call again - cache hit should return same result + assert descendants(fwd, 0, cache) == {1, 2, 3} + assert descendants(fwd, 1, cache) == {3} + + +def test_ancestors(): + fwd, rev, _ = build_adjacency(4, [(0, 1), (0, 2), (1, 3), (2, 3)]) + assert ancestors(rev, 3) == {0, 1, 2} + assert ancestors(rev, 0) == set() + + +def test_invert_adj(): + fwd, _, _ = build_adjacency(4, [(0, 1), (0, 2), (1, 3)]) + rev = invert_adj(fwd) + assert rev == [[], [0], [0], [1]] + + +def test_root_nodes(): + _, _, ind = build_adjacency(4, [(0, 2), (1, 2), (2, 3)]) + assert root_nodes(ind) == [0, 1] + + +def test_longest_path(): + # 0(w=1) -> 1(w=2) -> 3(w=4) + # 0(w=1) -> 2(w=10) -> 3(w=4) + fwd, _, ind = build_adjacency(4, [(0, 1), (0, 2), (1, 3), (2, 3)]) + topo = kahn_topological_order(4, fwd, ind) + lp = longest_path_length(fwd, topo, [1.0, 2.0, 10.0, 4.0]) + # longest from 0: 0 -> 2 -> 3 = 1+10+4 = 15 + assert lp[0] == 15.0 + assert lp[3] == 4.0 + assert lp[2] == 14.0 + + +def test_descendants_deep_chain_no_recursion_limit(): + # Chain of 2000 nodes -- would blow prototype's recursion limit + N = 2000 + edges = [(i, i + 1) for i in range(N - 1)] + fwd, _, _ = build_adjacency(N, edges) + result = descendants(fwd, 0) + assert len(result) == N - 1 diff --git a/MC/workflow_runner/o2dpg_runner/tests/test_resources.py b/MC/workflow_runner/o2dpg_runner/tests/test_resources.py new file mode 100644 index 000000000..2a723fc87 --- /dev/null +++ b/MC/workflow_runner/o2dpg_runner/tests/test_resources.py @@ -0,0 +1,137 @@ +import pytest + +from o2dpg_runner.resources import ( + ResourceManager, ResourceLimitExceeded, Semaphore, TaskResources, + ResourceBoundaries, +) + + +def _make_rm(cpu=8.0, mem=16000.0, **kw): + return ResourceManager(cpu_limit=cpu, mem_limit=mem, **kw) + + +def test_add_task_within_limits(): + rm = _make_rm() + rm.add_task("a", None, cpu=2, cpu_relative=1, mem=1000) + assert len(rm.resources) == 1 + assert rm.resources[0].name == "a" + + +def test_add_task_exceeds_limits_raises(): + rm = _make_rm(cpu=4) + with pytest.raises(ResourceLimitExceeded): + rm.add_task("big", None, cpu=8, cpu_relative=1, mem=100) + + +def test_is_within_limits_checks_mem_correctly(): + """Regression: prototype compared CPU to mem_limit by mistake.""" + rm = _make_rm(cpu=100, mem=1000, optimistic_resources=True) + # CPU under limit, MEM over limit -> should be caught + res = rm.add_task("t", None, cpu=1, cpu_relative=1, mem=5000) + # After add_task's limit_resources, mem is capped + assert res.mem_assigned == 1000 + + +def test_book_and_unbook_default(): + rm = _make_rm() + rm.add_task("a", None, 2, 1, 1000) + res = rm.resources[0] + res.nice_value = rm.nice_default + rm.book(0, rm.nice_default) + assert rm.n_procs == 1 + assert rm.cpu_booked == 2 + assert rm.mem_booked == 1000 + rm.unbook(0) + assert rm.n_procs == 0 + assert rm.cpu_booked == 0 + assert rm.mem_booked == 0 + + +def test_book_and_unbook_backfill(): + rm = _make_rm() + rm.add_task("a", None, 2, 1, 1000) + res = rm.resources[0] + res.nice_value = rm.nice_backfill + rm.book(0, rm.nice_backfill) + assert rm.n_procs_backfill == 1 + assert rm.n_procs == 0 + rm.unbook(0) + assert rm.n_procs_backfill == 0 + + +def test_fits_default(): + rm = _make_rm(cpu=4, mem=4000) + rm.add_task("a", None, 2, 1, 1000) + rm.add_task("b", None, 3, 1, 1000) + rm.resources[0].nice_value = rm.nice_default + rm.book(0, rm.nice_default) + # now 2 cpu / 1000 mem booked; b (3 cpu) doesn't fit + assert not rm.fits_default(rm.resources[1]) + + +def test_fits_backfill_rejects_too_big(): + rm = _make_rm(cpu=4, mem=4000) + rm.add_task("big", None, 4, 1, 1000) # equals cpu_limit -> 100% > 90% + rm.resources[0].limit_resources() # ensure assigned=cpu_limit + rm.resources[0].cpu_assigned = 4 # at 90% threshold + assert not rm.fits_backfill(rm.resources[0]) + + +def test_semaphore_blocks_duplicate(): + rm = _make_rm() + rm.add_task("a", None, 1, 1, 100, semaphore_string="S") + rm.add_task("b", None, 1, 1, 100, semaphore_string="S") + assert rm.resources[0].semaphore is rm.resources[1].semaphore + rm.resources[0].nice_value = rm.nice_default + rm.book(0, rm.nice_default) + assert not rm.can_be_submitted_at_all(rm.resources[1]) + rm.unbook(0) + assert rm.can_be_submitted_at_all(rm.resources[1]) + + +def test_related_tasks_share_bucket(): + rm = _make_rm() + rm.add_task("sgnsim_1", "sgnsim", 2, 1, 1000) + rm.add_task("sgnsim_2", "sgnsim", 2, 1, 1000) + assert rm.resources[0].related_tasks is rm.resources[1].related_tasks + + +def test_dynamic_sampling_propagates(): + rm = ResourceManager(cpu_limit=8, mem_limit=16000, dynamic_resources=True) + rm.add_task("t_1", "t", 2, 1, 1000) + rm.add_task("t_2", "t", 2, 1, 1000) + # Feed monitor samples for t_1 + for i in range(5): + rm.add_monitored(0, i * 1.0, cpu_fraction=1.5, mem_mb=800.0) + rm.resources[0].nice_value = rm.nice_default + rm.book(0, rm.nice_default) + rm.unbook(0) # this triggers sampling + propagation + # t_2 should now have an adjusted assignment based on observed sample + assert rm.resources[1].cpu_assigned > 0 + # The sampled CPU was ~1.5; t_2's cpu_assigned should reflect that ballpark + assert rm.resources[1].mem_assigned == pytest.approx(800.0, abs=1) + + +def test_a_zero_cpu_sample_leaves_siblings_alone(): + """A task seen only through its psutil baseline reads 0.0 cores. Passing + that on would make every sibling look free and admit them all at once.""" + manager = ResourceManager(cpu_limit=8, mem_limit=16000, dynamic_resources=True) + manager.add_task("t_1", "t", 2, 1, 1000) + manager.add_task("t_2", "t", 2, 1, 1000) + for i in range(5): + manager.add_monitored(0, i * 1.0, cpu_fraction=0.0, mem_mb=800.0) + manager.resources[0].nice_value = manager.nice_default + manager.book(0, manager.nice_default) + manager.unbook(0) + assert manager.resources[1].cpu_assigned == pytest.approx(2.0) + + +def test_at_proc_cap(): + rm = _make_rm() + rm.procs_parallel_max = 2 + assert not rm.at_proc_cap() + for i in range(2): + rm.add_task(f"t{i}", None, 1, 1, 100) + rm.resources[i].nice_value = rm.nice_default + rm.book(i, rm.nice_default) + assert rm.at_proc_cap() diff --git a/MC/workflow_runner/o2dpg_runner/tests/test_scheduler.py b/MC/workflow_runner/o2dpg_runner/tests/test_scheduler.py new file mode 100644 index 000000000..720f87ae4 --- /dev/null +++ b/MC/workflow_runner/o2dpg_runner/tests/test_scheduler.py @@ -0,0 +1,179 @@ +import pytest + +from o2dpg_runner.resources import ResourceManager +from o2dpg_runner.scheduler import ( + TimeframeFirstPolicy, CriticalPathPolicy, BestFitBackfillPolicy, get_policy, +) +from o2dpg_runner.scheduler.base import SchedulerState + + +def _setup(n=4, cpu_limit=8.0, mem_limit=16000.0): + """Create an RM with n tasks: all 2 cpu / 1000 mem.""" + rm = ResourceManager(cpu_limit=cpu_limit, mem_limit=mem_limit, n_backfill_max=1) + for i in range(n): + rm.add_task(f"t{i}", None, cpu=2, cpu_relative=1, mem=1000) + return rm + + +def _make_state(n, tf=None, desc=None, cp=None): + tf = tf if tf is not None else [0] * n + desc = desc if desc is not None else [0] * n + cp = cp if cp is not None else [1.0] * n + return SchedulerState( + timeframe_of=tf, + descendants_count=desc, + critical_path=cp, + task_cpu=[2.0] * n, + task_mem=[1000.0] * n, + timeframe_weight=[(tf[i], desc[i]) for i in range(n)], + ) + + +def test_timeframe_policy_order(): + s = _make_state(4, tf=[1, 0, 1, 0], desc=[5, 2, 10, 1]) + p = TimeframeFirstPolicy() + # timeframe 0 first, then within tf more descendants first + order = p.order([0, 1, 2, 3], s) + # tf 0: tasks 1 (desc=2) and 3 (desc=1) -> 1 before 3 + # tf 1: tasks 0 (desc=5) and 2 (desc=10) -> 2 before 0 + assert order == [1, 3, 2, 0] + + +def test_timeframe_policy_submit_default_fits_all(): + rm = _setup(n=4) # 4 * 2cpu = 8 = cpu_limit + s = _make_state(4) + p = TimeframeFirstPolicy() + picked = _drain(p.pick_submittable([0, 1, 2, 3], rm), rm) + assert len(picked) == 4 + assert all(nice == rm.nice_default for _, nice in picked) + + +def _drain(picks, rm): + """Drain a pick_submittable generator, booking each pick as the + executor would (so subsequent fits_default checks see it).""" + out = [] + for tid, nice in picks: + rm.book(tid, nice) + out.append((tid, nice)) + return out + + +def test_timeframe_policy_should_break_legacy(): + """Prototype behavior: the first non-fitting task breaks the default + pass entirely; following tasks only get a shot in the backfill pass + (which itself may be blocked once an earlier task consumed its budget).""" + rm = ResourceManager(cpu_limit=4, mem_limit=16000, n_backfill_max=4) + rm.add_task("small1", None, cpu=2, cpu_relative=1, mem=100) + rm.add_task("big", None, cpu=3, cpu_relative=1, mem=100) + rm.add_task("small2", None, cpu=2, cpu_relative=1, mem=100) + p = TimeframeFirstPolicy(drop_should_break=False) + picks = _drain(p.pick_submittable([0, 1, 2], rm), rm) + nice_for = {tid: n for tid, n in picks} + # small1 fits default + assert nice_for[0] == rm.nice_default + # small2, which WOULD have fit default (2+2=4), did not go default -- + # that's the legacy bug/feature. It's either backfill or not scheduled. + assert nice_for.get(2) != rm.nice_default + + +def test_timeframe_policy_drop_unblocks_small2_at_default(tmp_path=None): + """Contrast: with drop_should_break=True, small2 gets default priority.""" + rm = ResourceManager(cpu_limit=4, mem_limit=16000, n_backfill_max=4) + rm.add_task("small1", None, cpu=2, cpu_relative=1, mem=100) + rm.add_task("big", None, cpu=3, cpu_relative=1, mem=100) + rm.add_task("small2", None, cpu=2, cpu_relative=1, mem=100) + p = TimeframeFirstPolicy(drop_should_break=True) + picks = _drain(p.pick_submittable([0, 1, 2], rm), rm) + nice_for = {tid: n for tid, n in picks} + # small1 and small2 both at default -- the poster's headline result + assert nice_for[0] == rm.nice_default + assert nice_for[2] == rm.nice_default + + +def test_timeframe_policy_drop_should_break(): + """With drop_should_break=True, light tasks can slip past heavy ones + in the default pass.""" + rm = ResourceManager(cpu_limit=4, mem_limit=16000, n_backfill_max=4) + rm.add_task("small1", None, cpu=2, cpu_relative=1, mem=100) + rm.add_task("big", None, cpu=3, cpu_relative=1, mem=100) + rm.add_task("small2", None, cpu=2, cpu_relative=1, mem=100) + s = _make_state(3) + p = TimeframeFirstPolicy(drop_should_break=True) + picks = _drain(p.pick_submittable([0, 1, 2], rm), rm) + nice_for = {tid: n for tid, n in picks} + # small1 fits default; big doesn't (2+3=5>4); small2 fits default (2+2=4) + assert nice_for.get(0) == rm.nice_default + assert nice_for.get(2) == rm.nice_default + + +def test_critical_path_policy_order(): + # node 2 has the highest critical path -> should come first + s = _make_state(3, cp=[5.0, 3.0, 10.0]) + p = CriticalPathPolicy() + assert p.order([0, 1, 2], s) == [2, 0, 1] + + +def test_best_fit_policy_fills_budget(): + # cpu_limit=10; tasks: 3,3,4 -> best-fit should pack all three + rm = ResourceManager(cpu_limit=10, mem_limit=16000, n_backfill_max=1) + rm.add_task("a", None, cpu=3, cpu_relative=1, mem=100) + rm.add_task("b", None, cpu=3, cpu_relative=1, mem=100) + rm.add_task("c", None, cpu=4, cpu_relative=1, mem=100) + s = _make_state(3, cp=[1.0, 1.0, 1.0], desc=[0, 0, 0]) + p = BestFitBackfillPolicy() + ordered = p.order([0, 1, 2], s) + picks = _drain(p.pick_submittable(ordered, rm), rm) + assert len(picks) == 3 + assert all(n == rm.nice_default for _, n in picks) + # verify that with overfilled budget, best-fit DOES stop + rm2 = ResourceManager(cpu_limit=5, mem_limit=16000, n_backfill_max=0) + rm2.add_task("a", None, cpu=3, cpu_relative=1, mem=100) + rm2.add_task("b", None, cpu=3, cpu_relative=1, mem=100) + ordered2 = p.order([0, 1], _make_state(2)) + picks2 = _drain(p.pick_submittable(ordered2, rm2), rm2) + assert len(picks2) == 1 # only one fits the cpu=5 budget + + +def test_get_policy_names(): + assert isinstance(get_policy("timeframe"), TimeframeFirstPolicy) + assert isinstance(get_policy("critical-path"), CriticalPathPolicy) + assert isinstance(get_policy("best-fit"), BestFitBackfillPolicy) + with pytest.raises(ValueError): + get_policy("nonsense") + + +def test_semaphore_prevents_concurrent_submit(): + rm = ResourceManager(cpu_limit=8, mem_limit=16000, n_backfill_max=1) + rm.add_task("a", None, 1, 1, 100, semaphore_string="S") + rm.add_task("b", None, 1, 1, 100, semaphore_string="S") + s = _make_state(2) + # book a directly + rm.resources[0].nice_value = rm.nice_default + rm.book(0, rm.nice_default) + p = TimeframeFirstPolicy() + picks = list(p.pick_submittable([1], rm)) + assert picks == [] + + +def test_n_backfill_cap(): + rm = ResourceManager(cpu_limit=2, mem_limit=16000, n_backfill_max=1) + rm.add_task("a", None, 1, 1, 100) + rm.add_task("b", None, 1, 1, 100) + rm.add_task("c", None, 1, 1, 100) + # fill default with a+b + for i in (0, 1): + rm.resources[i].nice_value = rm.nice_default + rm.book(i, rm.nice_default) + assert rm.cpu_booked == 2 + # c doesn't fit default; backfill allowed? n_backfill_max=1 + p = TimeframeFirstPolicy() + picks = list(p.pick_submittable([2], rm)) + # one backfill slot, c gets it + assert len(picks) == 1 + assert picks[0][1] == rm.nice_backfill + rm.resources[2].nice_value = rm.nice_backfill + rm.book(2, rm.nice_backfill) + # next candidate tries to backfill but cap reached + rm.add_task("d", None, 1, 1, 100) + picks2 = list(p.pick_submittable([3], rm)) + assert picks2 == [] diff --git a/MC/workflow_runner/o2dpg_runner/tests/test_simulator.py b/MC/workflow_runner/o2dpg_runner/tests/test_simulator.py new file mode 100644 index 000000000..727bb394f --- /dev/null +++ b/MC/workflow_runner/o2dpg_runner/tests/test_simulator.py @@ -0,0 +1,131 @@ +import os +import sys + +import pytest + +from o2dpg_runner.workflow import build_workflow + +BIN_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) +if BIN_DIR not in sys.path: + sys.path.insert(0, BIN_DIR) + +import o2dpg_schedule_simulator as sim + + +def _wf(stages): + return build_workflow({"stages": stages}, ["*"], []) + + +def _task(name, cpu, mem=100.0, walltime=1.0, needs=None, timeframe=1): + return { + "name": name, + "needs": list(needs or []), + "cmd": "true", + "cwd": ".", + "timeframe": timeframe, + "labels": [], + "resources": {"cpu": cpu, "mem": mem, "walltime": walltime}, + } + + +def test_simulator_uses_amdahl_override_in_critical_path_state(): + wf = _wf([ + _task("a", cpu=1, walltime=1.0), + _task("b", cpu=1, walltime=5.0), + ]) + model = sim.AmdahlModel( + t_serial=20.0, + t_parallel_tot=0.0, + n_ref=1, + cpu_mean_ref=1.0, + min_workers=1, + max_workers=1, + ) + result = sim.simulate( + wf, + "critical-path", + cpu_limit=8.0, + mem_limit=1000.0, + amdahl_models={"a": model}, + worker_assignment={"a": 1}, + maxjobs=1, + ) + assert result.tasks[0].name == "a" + + +def test_simulator_keeps_tid_mapping_when_task_exceeds_limits(): + wf = _wf([ + _task("too_big", cpu=20, walltime=3.0), + _task("ok", cpu=1, walltime=2.0), + ]) + result = sim.simulate( + wf, + "timeframe", + cpu_limit=4.0, + mem_limit=1000.0, + maxjobs=1, + ) + assert [t.name for t in result.tasks] == ["ok"] + assert result.deadlocked_tids == [0] + + +def test_simulator_backfill_slowdown_marks_and_slows_backfill_tasks(): + wf = _wf([ + _task("small1", cpu=2, walltime=3.0), + _task("big", cpu=3, walltime=8.0), + _task("small2", cpu=2, walltime=3.0), + ]) + result = sim.simulate( + wf, + "timeframe", + cpu_limit=4.0, + mem_limit=1000.0, + backfill_model="slowdown", + n_backfill=1, + backfill_slowdown_factor=1.25, + ) + by_name = {t.name: t for t in result.tasks} + assert "big" in by_name + assert by_name["big"].start == pytest.approx(0.0) + compute_wt = 8.0 * 1.25 + assert by_name["big"].walltime == pytest.approx(compute_wt + 0.1) + # cpu is the average over the whole slot, so the idle overhead dilutes it; + # cpu * walltime is the CPU-seconds actually spent + assert by_name["big"].cpu == pytest.approx( + 3.0 / 1.25 * compute_wt / (compute_wt + 0.1)) + assert by_name["big"].cpu_booked == pytest.approx(3.0) + assert result.cpu_utilization(4.0) <= 1.0 + + +def test_simulator_holefill_uses_only_foreground_hole(): + wf = _wf([ + _task("fg", cpu=2, walltime=4.0), + _task("bf", cpu=3, walltime=2.0), + ]) + result = sim.simulate( + wf, + "timeframe", + cpu_limit=4.0, + mem_limit=1000.0, + backfill_model="holefill", + n_backfill=1, + task_overhead=0.0, + ) + by_name = {t.name: t for t in result.tasks} + assert by_name["fg"].walltime == pytest.approx(4.0) + assert by_name["bf"].walltime == pytest.approx(3.0) + assert by_name["bf"].cpu == pytest.approx(2.0) + assert by_name["bf"].cpu_booked == pytest.approx(3.0) + assert result.cpu_utilization(4.0) == pytest.approx(0.875) + + +def test_amdahl_model_rejects_negative_components(): + with pytest.raises(ValueError): + sim.AmdahlModel.from_dict( + { + "t_serial": -1.0, + "t_parallel_tot": 5.0, + "n_ref": 4, + "cpu_mean_ref": 4.0, + } + ) diff --git a/MC/workflow_runner/o2dpg_runner/tests/test_workflow.py b/MC/workflow_runner/o2dpg_runner/tests/test_workflow.py new file mode 100644 index 000000000..db912673b --- /dev/null +++ b/MC/workflow_runner/o2dpg_runner/tests/test_workflow.py @@ -0,0 +1,117 @@ +import json +import os + +import pytest + +from o2dpg_runner.workflow import ( + load_json, extract_global_init, filter_workflow, build_workflow, + update_resource_estimates, +) + +FIXTURE = os.path.join(os.path.dirname(__file__), "fixtures", "tiny_workflow.json") + + +def _load(): + return load_json(FIXTURE) + + +def test_extract_global_init(): + spec = _load() + env, cmd = extract_global_init(spec) + assert env == {"FOO": "bar"} + assert cmd is None # cmd was "NO-COMMAND" + # __global_init_task__ removed from stages + assert spec["stages"][0]["name"] != "__global_init_task__" + assert spec["stages"][0]["name"] == "bkg" + + +def test_filter_workflow_all(): + spec = _load() + extract_global_init(spec) + filtered, targets = filter_workflow(spec, ["*"], []) + assert targets == [] # "*" means no target list + assert len(filtered["stages"]) == len(spec["stages"]) + + +def test_filter_workflow_by_target(): + spec = _load() + extract_global_init(spec) + filtered, targets = filter_workflow(spec, ["digi_1"], []) + names = [t["name"] for t in filtered["stages"]] + # digi_1 + its needs bkg, sgnsim_1 + assert "digi_1" in names + assert "sgnsim_1" in names + assert "bkg" in names + assert "reco_1" not in names + assert "aod" not in names + assert targets == ["digi_1"] + + +def test_filter_workflow_by_label(): + spec = _load() + extract_global_init(spec) + filtered, targets = filter_workflow(spec, ["*"], ["QC"]) + # "*" with label filter narrows, and pulls in all deps + names = [t["name"] for t in filtered["stages"]] + assert "qc_1" in names + assert "qc_2" in names + # deps pulled in + for need in ("reco_1", "reco_2", "digi_1", "digi_2", "sgnsim_1", "sgnsim_2", "bkg"): + assert need in names + + +def test_filter_workflow_regex_target(): + spec = _load() + extract_global_init(spec) + filtered, targets = filter_workflow(spec, ["^qc_.*"], []) + names = [t["name"] for t in filtered["stages"]] + assert set(targets) == {"qc_1", "qc_2"} + assert "aod" not in names + + +def test_filter_does_not_alias_input(): + """Regression: the prototype did `transformedworkflowspec = workflowspec` + and then mutated .stages, aliasing the caller's dict. Ensure we don't.""" + spec = _load() + extract_global_init(spec) + original_len = len(spec["stages"]) + filtered, _ = filter_workflow(spec, ["qc_1"], []) + # filtered should be smaller, spec unchanged + assert len(filtered["stages"]) < original_len + assert len(spec["stages"]) == original_len + + +def test_build_workflow_end_to_end(): + raw = _load() + wf = build_workflow(raw, ["*"], []) + # 10 tasks after removing global init + assert wf.n_tasks() == 10 + # indegrees: bkg has 0; aod has 2 (reco_1, reco_2) + bkg_tid = wf.tid("bkg") + aod_tid = wf.tid("aod") + assert wf.indegree[bkg_tid] == 0 + assert wf.indegree[aod_tid] == 2 + # timeframes: -1, 1, 2 + assert wf.timeframes == {-1, 1, 2} + + +def test_update_resource_estimates(tmp_path): + raw = _load() + wf = build_workflow(raw, ["*"], []) + # synthesize a learned-estimates JSON keyed by "global" task name + est = { + "sgnsim": {"pss": {"max": 3000}, "cpu": {"mean": 3.5}}, + "digi": {"pss": {"max": 1200}, "cpu": {"mean": 1.8}}, + } + p = tmp_path / "res.json" + p.write_text(json.dumps(est)) + update_resource_estimates(wf, str(p)) + + for name in ("sgnsim_1", "sgnsim_2"): + t = wf.stages[wf.tid(name)] + assert t["resources"]["mem"] == 3000 + assert t["resources"]["cpu"] == 3.5 + for name in ("digi_1", "digi_2"): + t = wf.stages[wf.tid(name)] + assert t["resources"]["mem"] == 1200 + assert t["resources"]["cpu"] == 1.8 diff --git a/MC/workflow_runner/o2dpg_runner/workflow.py b/MC/workflow_runner/o2dpg_runner/workflow.py new file mode 100644 index 000000000..7c2322f8d --- /dev/null +++ b/MC/workflow_runner/o2dpg_runner/workflow.py @@ -0,0 +1,411 @@ +"""Workflow loading, filtering, and DAG construction. + +The workflow JSON schema is unchanged from the original. Key fields per +stage: + - name: str + - needs: list[str] (names of upstream stages) + - cmd: str + - cwd: str + - timeframe: int (-1 for global stages) + - labels: list[str] + - resources: {cpu, mem, relative_cpu} + - semaphore: str (optional) + - retry_count: int (optional) + - alternative_alienv_package: str (optional) + - env: dict (optional) + +One stage may be the synthetic ``__global_init_task__`` at index 0, +holding global env and an optional init cmd; it is stripped from the +DAG during loading. +""" + +from __future__ import annotations + +import copy +import json +import logging +import math +import re +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional, Set, Tuple + +from .graph import build_adjacency + +log = logging.getLogger(__name__) + + +@dataclass +class Workflow: + """In-memory representation of a workflow after filtering. + + Tasks are stored as a list of dicts (the raw JSON objects) and indexed + by integer ``tid``. The forward/reverse adjacency and derived quantities + are computed once at construction time. + """ + stages: List[Dict[str, Any]] + global_env: Dict[str, str] = field(default_factory=dict) + global_init_cmd: Optional[str] = None + full_target_names: List[str] = field(default_factory=list) + + # Derived + name_to_id: Dict[str, int] = field(default_factory=dict) + id_to_name: List[str] = field(default_factory=list) + forward_adj: List[List[int]] = field(default_factory=list) + reverse_adj: List[List[int]] = field(default_factory=list) + indegree: List[int] = field(default_factory=list) + timeframes: Set[int] = field(default_factory=set) + + def __post_init__(self): + self._rebuild_indices() + + def _rebuild_indices(self): + self.name_to_id = {s["name"]: i for i, s in enumerate(self.stages)} + self.id_to_name = [s["name"] for s in self.stages] + edges: List[Tuple[int, int]] = [] + for i, s in enumerate(self.stages): + for n in s.get("needs", []): + if n in self.name_to_id: + edges.append((self.name_to_id[n], i)) + self.forward_adj, self.reverse_adj, self.indegree = build_adjacency( + len(self.stages), edges + ) + self.timeframes = {s.get("timeframe", -1) for s in self.stages} + + def tid(self, name: str) -> int: + return self.name_to_id[name] + + def name(self, tid: int) -> str: + return self.id_to_name[tid] + + def n_tasks(self) -> int: + return len(self.stages) + + +def load_json(path: str) -> Dict[str, Any]: + with open(path) as fp: + return json.load(fp) + + +def extract_global_init(raw_spec: Dict[str, Any]) -> Tuple[Dict[str, str], Optional[str]]: + """Pull out the synthetic __global_init_task__ if present. + + Mutates raw_spec['stages'] in place (removes the init stage). + """ + env: Dict[str, str] = {} + init_cmd: Optional[str] = None + stages = raw_spec.get("stages", []) + if stages and stages[0].get("name") == "__global_init_task__": + init = stages[0] + env_in = init.get("env") + if env_in: + env = {k: str(v) for k, v in env_in.items()} + cmd = init.get("cmd") + if cmd and cmd != "NO-COMMAND": + init_cmd = cmd + del stages[0] + return env, init_cmd + + +def filter_workflow( + raw_spec: Dict[str, Any], + targets: List[str], + target_labels: List[str], +) -> Tuple[Dict[str, Any], List[str]]: + """Filter the raw spec down to tasks matching target selectors. + + Returns (new_spec, full_target_names). When no filter is requested, + returns (raw_spec, []). The returned spec is always a fresh top-level + dict (no aliasing bug like in the prototype), but the per-stage dicts + are shared. + """ + stages = raw_spec.get("stages", []) + if not targets: + return {**raw_spec, "stages": list(stages)}, [] + if not target_labels and len(targets) == 1 and targets[0] == "*": + return {**raw_spec, "stages": list(stages)}, [] + + name_to_idx = {t["name"]: i for i, t in enumerate(stages)} + + def task_matches(name: str) -> bool: + for f in targets: + if f == "*": + return True + if re.match(f, name) is not None: + return True + return False + + def task_matches_labels(t: Dict[str, Any]) -> bool: + if not target_labels: + return True + for lbl in t.get("labels", []): + if lbl in target_labels: + return True + return False + + # Memoized canBeDone using iterative traversal. + ok_cache: Dict[str, bool] = {} + + def can_be_done(name: str) -> bool: + if name in ok_cache: + return ok_cache[name] + idx = name_to_idx.get(name) + if idx is None: + ok_cache[name] = False + return False + # iterative post-order DFS + order: List[str] = [] + seen: Set[str] = {name} + stack: List[Tuple[str, int]] = [(name, 0)] + while stack: + cur, ci = stack[-1] + needs = stages[name_to_idx[cur]].get("needs", []) if cur in name_to_idx else [] + if ci < len(needs): + stack[-1] = (cur, ci + 1) + child = needs[ci] + if child not in seen and child not in ok_cache: + if child not in name_to_idx: + ok_cache[child] = False + else: + seen.add(child) + stack.append((child, 0)) + else: + stack.pop() + order.append(cur) + for cur in order: + if cur in ok_cache: + continue + idx2 = name_to_idx.get(cur) + if idx2 is None: + ok_cache[cur] = False + continue + ok = all(ok_cache.get(r, False) for r in stages[idx2].get("needs", [])) + ok_cache[cur] = ok + if not ok: + log.info("Disabling target %s due to unsatisfied requirements", cur) + return ok_cache[name] + + full_target_list = [ + t for t in stages + if task_matches(t["name"]) and task_matches_labels(t) and can_be_done(t["name"]) + ] + full_target_names = [t["name"] for t in full_target_list] + + # Collect all upstream requirements (iterative, deduped). + needed: Set[str] = set(full_target_names) + stack2 = list(full_target_names) + while stack2: + cur = stack2.pop() + idx = name_to_idx.get(cur) + if idx is None: + continue + for r in stages[idx].get("needs", []): + if r not in needed: + needed.add(r) + stack2.append(r) + + new_stages = [t for t in stages if t["name"] in needed] + new_spec = {**raw_spec, "stages": new_stages} + return new_spec, full_target_names + + +def build_workflow( + raw_spec: Dict[str, Any], + targets: List[str], + target_labels: List[str], +) -> Workflow: + """End-to-end: strip global-init, filter, build DAG.""" + # Operate on a shallow copy at top level so we don't mutate caller's dict. + spec = {**raw_spec, "stages": list(raw_spec.get("stages", []))} + # extract_global_init mutates spec['stages'] + env, init_cmd = extract_global_init(spec) + filtered, target_names = filter_workflow(spec, targets, target_labels) + wf = Workflow( + stages=filtered["stages"], + global_env=env, + global_init_cmd=init_cmd, + full_target_names=target_names, + ) + return wf + + +def update_resource_estimates( + workflow: Workflow, + resource_json_path: str, + logger=None, +) -> None: + """Apply learned resource estimates from a JSON file. + + The JSON is produced by o2dpg_sim_metrics.py json-stat and is keyed on + the "global" task name (i.e. with the _ suffix stripped). + + MEM is taken from pss.max (peak proportional set size). + CPU is taken from cpu.mean (average cores used during the task). + + Note on relative_cpu: the workflow JSON carries a relative_cpu field + that historically scaled a "max" CPU estimate down to an "expected" + usage. When injecting *measured* cpu.mean values that scaling must NOT + be applied again — the measurement already reflects actual usage. + relative_cpu remains untouched and continues to be used by the dynamic- + resources sampler (resources.py) for sibling reassignment, which is + correct behaviour: the sampler scales a freshly-observed aggregate back + to an expected per-task assignment. + """ + _log = logger if logger is not None else log + _log.info("Applying learned resource estimates from: %s", resource_json_path) + + with open(resource_json_path) as fp: + resource_dict = json.load(fp) + + # Remove the metadata key so task lookup doesn't match it. + resource_dict.pop("count", None) + + n_stages = len(workflow.stages) + n_updated = 0 + missing_base_names: set = set() + + for task in workflow.stages: + tf = task.get("timeframe", -1) + name = task["name"] + global_name = "_".join(name.split("_")[:-1]) if tf >= 1 else name + + if global_name not in resource_dict: + missing_base_names.add(global_name) + continue + + new_res = resource_dict[global_name] + task_updated = False + + walltime = new_res.get("lifetime", {}).get("mean") + if walltime is not None: + # Store even when walltime=0 (sub-10ms tasks that GNU time rounds + # to zero). The simulator clamps to a 1ms minimum so zero is + # handled gracefully; the fallback (cpu * factor) is always worse. + task["resources"]["walltime"] = float(walltime) + _log.info(" WALLTIME %-40s %.3f s", name, float(walltime)) + task_updated = True + + new_mem = new_res.get("pss", {}).get("max") + if new_mem is not None: + old_mem = task["resources"]["mem"] + task["resources"]["mem"] = new_mem + _log.info(" MEM %-40s %.1f MB -> %.1f MB", name, float(old_mem), new_mem) + task_updated = True + + new_cpu = new_res.get("cpu", {}).get("mean") + if new_cpu is not None: + old_cpu = task["resources"]["cpu"] + uses_dynamic_workers = "O2DPG_DYNAMIC_NWORKER_OVERWRITE" in task.get("cmd", "") + + if uses_dynamic_workers: + # Round cpu.mean to the nearest integer worker count and use + # that value for BOTH the scheduler's cpu booking and the + # actual NWORKERS setting. This keeps them consistent: the + # task will run with n_workers processes each using ~1 core, + # so total cpu ≈ n_workers = what we book. + n_workers = max(1, round(new_cpu)) + task["resources"]["cpu"] = float(n_workers) + if not isinstance(task.get("env"), dict): + task["env"] = {} + task["env"]["O2DPG_DYNAMIC_NWORKER_OVERWRITE"] = str(n_workers) + _log.info( + " CPU+NWORKERS %-36s cpu.mean=%.2f -> %d workers, %.0f cores booked", + name, new_cpu, n_workers, float(n_workers), + ) + else: + # No dynamic worker override: book cpu.mean directly. + # Do NOT apply relative_cpu scaling — the measurement already + # reflects actual usage. + task["resources"]["cpu"] = new_cpu + _log.info(" CPU %-40s %.3f cores -> %.3f cores", + name, float(old_cpu), new_cpu) + task_updated = True + + if task_updated: + n_updated += 1 + + if missing_base_names: + _log.info(" No learned data for: %s", ", ".join(sorted(missing_base_names))) + _log.info( + "Resource update done: %d/%d task stages updated (%d base name(s) not in learned data).", + n_updated, n_stages, len(missing_base_names), + ) + + +def replicate_workflow_for_timeframes(raw_spec: Dict[str, Any], M: int) -> Dict[str, Any]: + """Return a synthetic M-timeframe workflow derived from *raw_spec*. + + The original workflow may have N timeframes. This function: + 1. Detects per-TF stages (``timeframe >= 1``) and uses the + lowest-numbered timeframe as the canonical template. + 2. Instantiates the template for TF=1..M. + 3. Updates global stage dependencies so they reference exactly TF=1..M. + + Works for both M < N (shrink) and M > N (expand). + The returned dict shares no mutable state with *raw_spec*. + """ + stages = raw_spec.get("stages", []) + per_tf = [s for s in stages if s.get("timeframe", -1) >= 1] + global_stgs = [s for s in stages if s.get("timeframe", -1) < 1] + + if not per_tf: + return raw_spec # no per-TF template structure detected + + original_tf_set = {s["timeframe"] for s in per_tf} + min_tf = min(original_tf_set) + template_stages = [s for s in per_tf if s.get("timeframe") == min_tf] + template_names = {s["name"] for s in template_stages} + all_per_tf_names = {s["name"] for s in per_tf} + + def _base(name: str, tf: int) -> str: + sfx = f"_{tf}" + return name[:-len(sfx)] if name.endswith(sfx) else name + + # Replicate per-TF tasks for i = 1 .. M. + new_per_tf: List[Dict[str, Any]] = [] + for i in range(1, M + 1): + for tmpl in template_stages: + s = copy.deepcopy(tmpl) + base = _base(s["name"], min_tf) + s["name"] = f"{base}_{i}" + s["timeframe"] = i + new_needs: List[str] = [] + for need in tmpl.get("needs", []): + if need in template_names: + new_needs.append(f"{_base(need, min_tf)}_{i}") + else: + new_needs.append(need) + s["needs"] = new_needs + new_per_tf.append(s) + + # Update global stages: replace all per-TF deps with the full 1..M set, + # expanding each unique base name exactly once (deduplicates cross-TF refs). + new_global: List[Dict[str, Any]] = [] + for gstage in global_stgs: + s = copy.deepcopy(gstage) + new_needs_g: List[str] = [] + seen: Set[str] = set() + expanded_bases: Set[str] = set() + for need in gstage.get("needs", []): + if need not in all_per_tf_names: + if need not in seen: + new_needs_g.append(need) + seen.add(need) + continue + # Identify the base name (strip whichever TF suffix this entry has). + base = need + for otf in original_tf_set: + if need.endswith(f"_{otf}"): + base = need[:-len(f"_{otf}")] + break + if base in expanded_bases: + continue + expanded_bases.add(base) + for j in range(1, M + 1): + n = f"{base}_{j}" + if n not in seen: + new_needs_g.append(n) + seen.add(n) + s["needs"] = new_needs_g + new_global.append(s) + + return {**raw_spec, "stages": new_per_tf + new_global} diff --git a/MC/workflow_runner/o2dpg_schedule_simulator.py b/MC/workflow_runner/o2dpg_schedule_simulator.py new file mode 100755 index 000000000..bb9e80819 --- /dev/null +++ b/MC/workflow_runner/o2dpg_schedule_simulator.py @@ -0,0 +1,1268 @@ +#!/usr/bin/env python3 +"""Discrete-event simulator for the O2DPG workflow scheduler. + +Loads a workflow (and optionally applies learned resources), then +simulates scheduling under one or more policies in microseconds — +no processes are spawned. + +Assumptions (matching a kernel-enforced hard CPU limit): + - No nice / backfill tier: all tasks are submitted at the default + nice level against a single hard CPU budget. + - Memory is also a hard limit. + - Tasks run for exactly their `resources.walltime` seconds (set by + --update-resources). If walltime is absent, cpu * --walltime-per-core + is used as a proxy. + - Task parallelism is limited only by cpu_limit and mem_limit, not + by --maxjobs (the simulator sets the process cap to infinity). + +Usage examples +-------------- +Compare all three policies with learned resources: + + o2dpg_schedule_simulator.py \\ + -f workflow.json \\ + --update-resources learned.json \\ + --cpu-limit 8 --mem-limit 16384 \\ + --policies timeframe critical-path best-fit + +Single policy, verbose per-task schedule: + + o2dpg_schedule_simulator.py \\ + -f workflow.json --update-resources learned.json \\ + --cpu-limit 8 --mem-limit 16384 \\ + --policies critical-path --verbose + +JSON output for downstream analysis: + + o2dpg_schedule_simulator.py ... --output sim.json +""" + +from __future__ import annotations + +import argparse +import copy +import json +import math +import os +import random +import statistics +import sys +from dataclasses import dataclass, field +from typing import Dict, List, Optional, Set, Tuple + +_here = os.path.dirname(os.path.abspath(__file__)) +if _here not in sys.path: + sys.path.insert(0, _here) + +from o2dpg_runner.workflow import ( + build_workflow, + load_json, + replicate_workflow_for_timeframes, + update_resource_estimates, +) +from o2dpg_runner.resources import ResourceManager, ResourceLimitExceeded +from o2dpg_runner.scheduler import get_policy +from o2dpg_runner.scheduler.base import SchedulerState +from o2dpg_runner.scheduler.timeframe import TimeframeFirstPolicy +from o2dpg_runner.graph import descendants, longest_path_length, kahn_topological_order + + +# --------------------------------------------------------------------------- +# Data classes +# --------------------------------------------------------------------------- + +@dataclass +class SimTask: + tid: int + name: str + start: float # wall seconds from t=0 + finish: float + cpu: float # effective average cores consumed over [start, finish] + cpu_booked: float # cores booked for scheduler admission + mem: float # MB booked + walltime: float # finish - start + + +@dataclass +class _RunningBackfill: + task: SimTask + nominal_work: float + remaining_work: float + launch_seq: int + overhead_until: Optional[float] = None + + +@dataclass +class SimResult: + policy: str + makespan: float # total wall seconds + tasks: List[SimTask] = field(default_factory=list) + deadlocked_tids: List[int] = field(default_factory=list) + + def cpu_utilization(self, cpu_limit: float) -> float: + """Mean CPU utilisation as a fraction of cpu_limit.""" + if self.makespan <= 0 or cpu_limit <= 0: + return 0.0 + total = sum(t.cpu * t.walltime for t in self.tasks) + return total / (self.makespan * cpu_limit) + + def peak_mem_mb(self) -> float: + """Peak concurrent memory usage in MB (sweep-line).""" + events: List[Tuple[float, float]] = [] + for t in self.tasks: + events.append((t.start, +t.mem)) + events.append((t.finish, -t.mem)) + events.sort() + peak = cur = 0.0 + for _, delta in events: + cur += delta + peak = max(peak, cur) + return peak + + def to_dict(self) -> dict: + return { + "policy": self.policy, + "makespan_s": round(self.makespan, 3), + "tasks": [ + { + "tid": t.tid, "name": t.name, + "start": round(t.start, 3), "finish": round(t.finish, 3), + "cpu": round(t.cpu, 3), "cpu_booked": round(t.cpu_booked, 3), + "mem": round(t.mem, 1), + "walltime": round(t.walltime, 3), + } + for t in sorted(self.tasks, key=lambda x: x.start) + ], + "deadlocked_tids": self.deadlocked_tids, + } + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _global_name(name: str) -> str: + toks = name.split("_") + if toks and toks[-1].isdigit() and len(toks) > 1: + return "_".join(toks[:-1]) + return name + + +def _task_walltime(task: dict, cpu_fallback_factor: float) -> float: + """Return walltime [s] for a task, falling back to cpu * factor.""" + wt = task.get("resources", {}).get("walltime") + if wt is not None: + try: + return max(1e-3, float(wt)) + except (TypeError, ValueError): + pass + cpu = float(task.get("resources", {}).get("cpu", 1.0)) + return max(1e-3, cpu * cpu_fallback_factor) + + +_LEARN_ALL: Set[str] = frozenset({"cpu", "mem", "lifetime"}) + + +def _apply_learned_fields(workflow, learned: Dict, fields: Set[str]) -> int: + """Patch workflow stages in-place with a subset of learned resource fields. + + *fields* is a subset of ``{"cpu", "mem", "lifetime"}``. Only the named + dimensions are written; the rest keep the values from workflow.json. + Returns the count of stages that received at least one update. + """ + n_updated = 0 + for task in workflow.stages: + tf = task.get("timeframe", -1) + name = task["name"] + gname = "_".join(name.split("_")[:-1]) if tf >= 1 else name + data = learned.get(gname) + if not isinstance(data, dict): + continue + updated = False + if "lifetime" in fields: + wt = data.get("lifetime", {}).get("mean") + if wt is not None: + task["resources"]["walltime"] = float(wt) + updated = True + if "mem" in fields: + mem = data.get("pss", {}).get("max") + if mem is not None: + task["resources"]["mem"] = float(mem) + updated = True + if "cpu" in fields: + cpu = data.get("cpu", {}).get("mean") + if cpu is not None: + task["resources"]["cpu"] = float(cpu) + updated = True + if updated: + n_updated += 1 + return n_updated + + +@dataclass +class AmdahlModel: + """Amdahl scaling model derived from a single measurement point. + + walltime(n) = t_serial + t_parallel_tot / n + + t_serial and t_parallel_tot are solved from: + walltime_ref = t_serial + t_parallel_tot / n_ref + cpu_mean_ref = (t_serial + t_parallel_tot) / walltime_ref + """ + t_serial: float + t_parallel_tot: float + n_ref: int + cpu_mean_ref: float + min_workers: int = 1 + max_workers: int = 1 + + def walltime(self, n: int) -> float: + return max(1e-3, self.t_serial + self.t_parallel_tot / max(1, n)) + + @property + def worker_range(self) -> List[int]: + return list(range(self.min_workers, self.max_workers + 1)) + + @classmethod + def from_dict(cls, d: dict) -> "AmdahlModel": + model = cls( + t_serial=float(d["t_serial"]), + t_parallel_tot=float(d["t_parallel_tot"]), + n_ref=int(d["n_ref"]), + cpu_mean_ref=float(d["cpu_mean_ref"]), + min_workers=int(d.get("min_workers", 1)), + max_workers=int(d.get("max_workers", d["n_ref"])), + ) + if model.t_serial < 0 or model.t_parallel_tot < 0: + raise ValueError("Amdahl model has negative serial/parallel component") + if model.n_ref < 1 or model.min_workers < 1 or model.max_workers < model.min_workers: + raise ValueError("Amdahl model has invalid worker bounds") + return model + + +def _sample_walltime(mean: float, std: float, rng: random.Random) -> float: + """Draw a walltime sample from a log-normal distribution. + + Log-normal is a natural model for walltime: always positive, right-skewed + (occasional slow outliers). When std=0 the mean is returned unchanged. + """ + if std <= 1e-9 or mean <= 1e-9: + return max(1e-3, mean) + cv = std / mean + sigma2 = math.log(1.0 + cv * cv) + mu = math.log(mean) - 0.5 * sigma2 + return max(1e-3, rng.lognormvariate(mu, math.sqrt(sigma2))) + + +def _build_rm( + workflow, + cpu_limit: float, + mem_limit: float, + cpu_overrides: Optional[Dict[int, float]] = None, + n_backfill_max: int = 0, + backfill_cpu_factor: float = 1.5, + backfill_mem_factor: float = 1.5, + maxjobs: int = 10_000, +) -> Tuple[ResourceManager, Set[int]]: + """Fresh ResourceManager with no backfill tier and unlimited job slots. + + *cpu_overrides* maps tid → cpu to override resources.cpu for specific + tasks (used by the worker-count optimizer). + """ + rm = ResourceManager( + cpu_limit=cpu_limit, + mem_limit=mem_limit, + procs_parallel_max=maxjobs, + n_backfill_max=n_backfill_max, + backfill_cpu_factor=backfill_cpu_factor, + backfill_mem_factor=backfill_mem_factor, + dynamic_resources=False, + optimistic_resources=True, + ) + impossible_tids: Set[int] = set() + for i, task in enumerate(workflow.stages): + rel = None + try: + rv = task["resources"].get("relative_cpu") + rel = float(rv) if rv is not None else None + except (TypeError, ValueError): + pass + cpu = float(task["resources"]["cpu"]) + if cpu_overrides and i in cpu_overrides: + cpu = cpu_overrides[i] + mem = float(task["resources"]["mem"]) + if cpu > cpu_limit or mem > mem_limit: + impossible_tids.add(i) + try: + rm.add_task( + name=task["name"], + related_name=_global_name(task["name"]), + cpu=cpu, + cpu_relative=rel, + mem=mem, + semaphore_string=task.get("semaphore"), + ) + except ResourceLimitExceeded as e: + print(f" WARNING: task {task['name']} exceeds limits and will never run: {e}", + file=sys.stderr) + if impossible_tids: + names = [workflow.stages[i]["name"] for i in sorted(impossible_tids)] + print( + " WARNING: tasks exceed the hard simulator limits and will remain unscheduled: " + f"{names[:5]}{'...' if len(names) > 5 else ''}", + file=sys.stderr, + ) + return rm, impossible_tids + + +def _build_state( + workflow, + cpu_fallback_factor: float, + cpu_overrides: Optional[Dict[int, float]] = None, + walltime_overrides: Optional[Dict[int, float]] = None, +) -> SchedulerState: + n = workflow.n_tasks() + desc_cache: Dict = {} + desc_counts = [len(descendants(workflow.forward_adj, tid, desc_cache)) + for tid in range(n)] + + timeframe_of = [t.get("timeframe", -1) for t in workflow.stages] + tf_weight = [(timeframe_of[i], desc_counts[i]) for i in range(n)] + + cpu = [float(t.get("resources", {}).get("cpu", 1.0)) for t in workflow.stages] + if cpu_overrides: + for tid, value in cpu_overrides.items(): + cpu[tid] = float(value) + mem = [float(t.get("resources", {}).get("mem", 0.0)) for t in workflow.stages] + walltime = [_task_walltime(t, cpu_fallback_factor) for t in workflow.stages] + if walltime_overrides: + for tid, value in walltime_overrides.items(): + walltime[tid] = float(value) + + has_walltime = any(t.get("resources", {}).get("walltime") for t in workflow.stages) + cp_weight = walltime if has_walltime else cpu + topo = kahn_topological_order(n, workflow.forward_adj, workflow.indegree) + cp = longest_path_length(workflow.forward_adj, topo, cp_weight) + + return SchedulerState( + timeframe_of=timeframe_of, + descendants_count=desc_counts, + critical_path=cp, + task_cpu=cpu, + task_mem=mem, + task_walltime=walltime, + timeframe_weight=tf_weight, + ) + + +# --------------------------------------------------------------------------- +# Core simulation +# --------------------------------------------------------------------------- + +def simulate( + workflow, + policy_name: str, + cpu_limit: float, + mem_limit: float, + cpu_fallback_factor: float = 10.0, + task_overhead: float = 0.1, + walltime_stds: Optional[Dict[str, float]] = None, + cv_fallback: float = 0.15, + rng: Optional[random.Random] = None, + amdahl_models: Optional[Dict[str, "AmdahlModel"]] = None, + worker_assignment: Optional[Dict[str, int]] = None, + backfill_model: str = "off", + n_backfill: int = 1, + backfill_cpu_factor: float = 1.5, + backfill_mem_factor: float = 1.5, + backfill_slowdown_factor: float = 1.15, + maxjobs: int = 10_000, +) -> SimResult: + """Run one discrete-event simulation; return SimResult. + + When *rng* is provided, each task's walltime is sampled from + log-normal(mean, std). The std is taken from *walltime_stds* when + available; otherwise *cv_fallback* × mean is used as a noise floor. + + When *amdahl_models* and *worker_assignment* are both provided, walltime + and cpu booking for scalable tasks are derived from the Amdahl model at + the assigned worker count rather than from the workflow resources. + """ + mean_walltimes = [_task_walltime(t, cpu_fallback_factor) for t in workflow.stages] + + # Apply Amdahl model overrides for scalable tasks. + cpu_overrides: Dict[int, float] = {} + walltime_overrides: Dict[int, float] = {} + if amdahl_models and worker_assignment: + for i, task in enumerate(workflow.stages): + base = _global_name(task["name"]) + model = amdahl_models.get(base) + n = worker_assignment.get(base) + if model is not None and n is not None: + mean_walltimes[i] = model.walltime(n) + walltime_overrides[i] = mean_walltimes[i] + cpu_overrides[i] = float(n) + + state = _build_state( + workflow, + cpu_fallback_factor, + cpu_overrides=cpu_overrides or None, + walltime_overrides=walltime_overrides or None, + ) + + # Sample walltimes for this simulation run. + if rng is not None: + walltimes = [] + for i, task in enumerate(workflow.stages): + mean_wt = mean_walltimes[i] + base = _global_name(task["name"]) + std_wt = (walltime_stds or {}).get(base, 0.0) + if std_wt <= 0 and cv_fallback > 0: + std_wt = mean_wt * cv_fallback + walltimes.append(_sample_walltime(mean_wt, std_wt, rng)) + else: + walltimes = mean_walltimes + + if policy_name == "timeframe": + policy = TimeframeFirstPolicy(drop_should_break=False) + else: + policy = get_policy(policy_name) + + use_backfill = backfill_model in ("structural", "slowdown", "holefill") + rm, impossible_tids = _build_rm( + workflow, + cpu_limit, + mem_limit, + cpu_overrides=cpu_overrides or None, + n_backfill_max=n_backfill if use_backfill else 0, + backfill_cpu_factor=backfill_cpu_factor, + backfill_mem_factor=backfill_mem_factor, + maxjobs=maxjobs, + ) + + n = workflow.n_tasks() + proc_status = ["ToDo"] * n + candidates: List[int] = [ + i for i in range(n) if workflow.indegree[i] == 0 and i not in impossible_tids + ] + finished: Set[int] = set() + running: List[Tuple[int, float]] = [] # (tid, finish_time) + result = SimResult(policy=policy_name, makespan=0.0) + t = 0.0 + + if backfill_model == "holefill": + running_fg: List[Tuple[int, float]] = [] + running_bf: Dict[int, _RunningBackfill] = {} + launch_seq = 0 + + for _guard in range(n * n + 1): + ordered = policy.order(candidates, state) + for tid, nice in policy.pick_submittable(ordered, rm): + if nice != rm.nice_default and rm.cpu_free_default() <= 1e-9: + continue + rm.book(tid, nice) + candidates.remove(tid) + proc_status[tid] = "Running" + res = rm.resources[tid] + if nice == rm.nice_default: + compute_wt = walltimes[tid] + wt = compute_wt + task_overhead + finish = t + wt + fg_cpu = res.cpu_assigned * (compute_wt / wt) if wt > 0 else 0.0 + task = SimTask( + tid=tid, + name=workflow.id_to_name[tid], + start=t, + finish=finish, + cpu=fg_cpu, + cpu_booked=res.cpu_assigned, + mem=res.mem_assigned, + walltime=wt, + ) + running_fg.append((tid, finish)) + result.tasks.append(task) + else: + launch_seq += 1 + task = SimTask( + tid=tid, + name=workflow.id_to_name[tid], + start=t, + finish=t, + cpu=0.0, + cpu_booked=res.cpu_assigned, + mem=res.mem_assigned, + walltime=0.0, + ) + nominal_work = res.cpu_assigned * walltimes[tid] + running_bf[tid] = _RunningBackfill( + task=task, + nominal_work=nominal_work, + remaining_work=nominal_work, + launch_seq=launch_seq, + ) + result.tasks.append(task) + + if not running_fg and not running_bf: + break + + next_fg = min((ft for _, ft in running_fg), default=float("inf")) + cpu_fg = sum( + next(task.cpu_booked for task in result.tasks if task.tid == tid) + for tid, _ in running_fg + ) + hole_cpu = max(0.0, cpu_limit - cpu_fg) + remaining_hole = hole_cpu + bf_alloc: Dict[int, float] = {} + for tid, rb in sorted(running_bf.items(), key=lambda item: item[1].launch_seq): + if rb.overhead_until is not None: + bf_alloc[tid] = 0.0 + continue + alloc = min(rb.task.cpu_booked, remaining_hole) + bf_alloc[tid] = alloc + remaining_hole -= alloc + + next_bf = float("inf") + for tid, rb in running_bf.items(): + if rb.overhead_until is not None: + next_bf = min(next_bf, rb.overhead_until) + else: + alloc = bf_alloc.get(tid, 0.0) + if alloc > 1e-12: + next_bf = min(next_bf, t + rb.remaining_work / alloc) + + next_t = min(next_fg, next_bf) + if not math.isfinite(next_t): + break + dt = max(0.0, next_t - t) + for tid, rb in running_bf.items(): + if rb.overhead_until is None: + alloc = bf_alloc.get(tid, 0.0) + if alloc > 0.0: + rb.remaining_work = max(0.0, rb.remaining_work - alloc * dt) + t = next_t + + new_running_fg: List[Tuple[int, float]] = [] + for tid, ft in running_fg: + if abs(ft - t) < 1e-9: + rm.unbook(tid) + proc_status[tid] = "Done" + finished.add(tid) + for succ in workflow.forward_adj[tid]: + if proc_status[succ] == "ToDo": + if succ in impossible_tids: + continue + if all(p in finished for p in workflow.reverse_adj[succ]): + candidates.append(succ) + else: + new_running_fg.append((tid, ft)) + running_fg = new_running_fg + + done_bf: List[int] = [] + for tid, rb in running_bf.items(): + if rb.overhead_until is not None: + if abs(rb.overhead_until - t) < 1e-9: + done_bf.append(tid) + continue + if rb.remaining_work <= 1e-9: + if task_overhead > 0: + rb.overhead_until = t + task_overhead + else: + done_bf.append(tid) + + for tid in done_bf: + rb = running_bf.pop(tid) + rb.task.finish = t + rb.task.walltime = max(1e-9, rb.task.finish - rb.task.start) + rb.task.cpu = rb.nominal_work / rb.task.walltime + rm.unbook(tid) + proc_status[tid] = "Done" + finished.add(tid) + for succ in workflow.forward_adj[tid]: + if proc_status[succ] == "ToDo": + if succ in impossible_tids: + continue + if all(p in finished for p in workflow.reverse_adj[succ]): + candidates.append(succ) + + if not candidates and not running_fg and not running_bf: + break + else: + print(f" WARNING [{policy_name}]: simulation hit guard limit — possible deadlock", + file=sys.stderr) + + result.makespan = t + result.deadlocked_tids = [ + i for i, s in enumerate(proc_status) if s == "ToDo" + ] + if result.deadlocked_tids: + names = [workflow.id_to_name[i] for i in result.deadlocked_tids] + print(f" WARNING [{policy_name}]: {len(names)} tasks never scheduled " + f"(resource limits too tight?): {names[:5]}{'...' if len(names) > 5 else ''}", + file=sys.stderr) + return result + + for _guard in range(n * n + 1): # at most n rounds to completion + # --- schedule everything that fits right now --- + ordered = policy.order(candidates, state) + for tid, nice in policy.pick_submittable(ordered, rm): + # book immediately so subsequent picks in this pass see the + # updated resource availability (mirrors executor behaviour) + rm.book(tid, nice) + slowdown = 1.0 + if backfill_model == "slowdown" and nice != rm.nice_default: + slowdown = backfill_slowdown_factor + compute_wt = walltimes[tid] * slowdown # time spent doing actual work + wt = compute_wt + task_overhead # total slot duration (incl. idle overhead) + finish = t + wt + running.append((tid, finish)) + candidates.remove(tid) + proc_status[tid] = "Running" + res = rm.resources[tid] + # Overhead is idle time (process start/stop, alienv load, I/O flush). + # Average CPU over the full slot = booked_cpu × compute_fraction only. + effective_cpu = res.cpu_assigned / slowdown * (compute_wt / wt) if wt > 0 else 0.0 + result.tasks.append(SimTask( + tid=tid, + name=workflow.id_to_name[tid], + start=t, + finish=finish, + cpu=effective_cpu, + cpu_booked=res.cpu_assigned, + mem=res.mem_assigned, + walltime=wt, + )) + + if not running: + break + + # --- advance to next task completion --- + next_t = min(ft for _, ft in running) + t = next_t + + # --- complete all tasks finishing at t (within float tolerance) --- + still_running: List[Tuple[int, float]] = [] + for tid, ft in running: + if abs(ft - t) < 1e-9: + rm.unbook(tid) + proc_status[tid] = "Done" + finished.add(tid) + for succ in workflow.forward_adj[tid]: + if proc_status[succ] == "ToDo": + if succ in impossible_tids: + continue + if all(p in finished for p in workflow.reverse_adj[succ]): + candidates.append(succ) + else: + still_running.append((tid, ft)) + running = still_running + + if not candidates and not running: + break + else: + print(f" WARNING [{policy_name}]: simulation hit guard limit — possible deadlock", + file=sys.stderr) + + result.makespan = t + result.deadlocked_tids = [ + i for i, s in enumerate(proc_status) if s == "ToDo" + ] + if result.deadlocked_tids: + names = [workflow.id_to_name[i] for i in result.deadlocked_tids] + print(f" WARNING [{policy_name}]: {len(names)} tasks never scheduled " + f"(resource limits too tight?): {names[:5]}{'...' if len(names) > 5 else ''}", + file=sys.stderr) + return result + + +# --------------------------------------------------------------------------- +# Presentation +# --------------------------------------------------------------------------- + +def _fmt_time(s: float) -> str: + return f"{s:.1f}s" + + +def print_summary( + results_by_policy: Dict[str, List[SimResult]], + cpu_limit: float, + n_samples: int, +) -> None: + w = 16 + stoch = n_samples > 1 + mk_hdr = f"{'Makespan (mean±std)':>22}" if stoch else f"{'Makespan':>10}" + cpu_hdr = f"{'CPU util (mean±std)':>20}" if stoch else f"{'CPU util':>9}" + mem_hdr = f"{'Peak mem (mean±std)':>22}" if stoch else f"{'Peak mem':>10}" + header = f"{'Policy':<{w}} {mk_hdr} {cpu_hdr} {mem_hdr} {'Tasks':>6}" + print() + print(header) + print("-" * len(header)) + for policy, runs in results_by_policy.items(): + makespans = [r.makespan for r in runs] + mean_mk = statistics.mean(makespans) + util_pct = statistics.mean(r.cpu_utilization(cpu_limit) * 100 for r in runs) + peak = statistics.mean(r.peak_mem_mb() for r in runs) + n_tasks = runs[0].tasks.__len__() if runs else 0 + if stoch: + std_mk = statistics.stdev(makespans) if len(makespans) > 1 else 0.0 + std_util = statistics.stdev(r.cpu_utilization(cpu_limit)*100 for r in runs) if len(runs)>1 else 0.0 + std_peak = statistics.stdev(r.peak_mem_mb() for r in runs) if len(runs)>1 else 0.0 + mk_str = f"{_fmt_time(mean_mk)} ± {_fmt_time(std_mk)}" + print(f"{policy:<{w}} {mk_str:>22} {util_pct:>7.1f}±{std_util:.1f}% {peak:>8.0f}±{std_peak:.0f}MB {n_tasks:>6}") + else: + print(f"{policy:<{w}} {_fmt_time(mean_mk):>10} {util_pct:>8.1f}% " + f"{peak:>9.0f}MB {n_tasks:>6}") + print() + + +def print_verbose(result: SimResult) -> None: + print(f"\n--- Schedule: {result.policy} ---") + prev_t = -1.0 + for task in sorted(result.tasks, key=lambda x: (x.start, x.name)): + if abs(task.start - prev_t) > 1e-9: + print(f" t={_fmt_time(task.start)}") + prev_t = task.start + print(f" START {task.name:<40} " + f"cpu={task.cpu:.1f}" + + (f" ({task.cpu_booked:.1f} booked)" if abs(task.cpu - task.cpu_booked) > 1e-9 else "") + + f" mem={task.mem:.0f}MB " + f"dur={_fmt_time(task.walltime)}") + print(f" Makespan: {_fmt_time(result.makespan)}") + + +def _print_sweep_table( + sweep_results: "List[Tuple]", # (M, results_by_policy, n_stages, worker_assignment|None) + cpu_limit: float, + n_samples: int, +) -> None: + """Print a compact table summarising a timeframe-sweep simulation run.""" + stoch = n_samples > 1 + + # Collect all scalable task names that appear in any worker assignment. + scalable_names: List[str] = [] + seen_names: "Set[str]" = set() + for _, _, _, wa in sweep_results: + if wa: + for name in sorted(wa): + if name not in seen_names: + scalable_names.append(name) + seen_names.add(name) + + # Build header with optional per-task worker columns. + worker_cols = " ".join(f"{n[:12]:>12}" for n in scalable_names) + cpu_col_hdr = f"{'CPU util(±std)':>14}" if stoch else f"{'CPU util':>9}" + hdr = (f" {'M':>4} {'Policy':<16} {'N tasks':>7} " + f"{'Makespan':>12} {cpu_col_hdr} {'Peak mem':>9}" + + (f" {worker_cols}" if scalable_names else "")) + cpu_col_w = 14 if stoch else 9 + print("\nTimeframe sweep results:") + print(hdr) + print(" " + "-" * (len(hdr) - 2)) + + for M, results_by_policy, n_stages, worker_assignment in sweep_results: + m_str = str(M) if M is not None else "orig" + for policy, runs in results_by_policy.items(): + makespans = [r.makespan for r in runs] + mean_mk = statistics.mean(makespans) + util_pct = statistics.mean(r.cpu_utilization(cpu_limit) * 100 for r in runs) + peak = statistics.mean(r.peak_mem_mb() for r in runs) + if stoch and len(runs) > 1: + std_mk = statistics.stdev(makespans) + mk_str = f"{_fmt_time(mean_mk)}±{_fmt_time(std_mk)}" + std_util = statistics.stdev(r.cpu_utilization(cpu_limit) * 100 for r in runs) + util_str = f"{util_pct:.1f}±{std_util:.1f}%" + else: + mk_str = _fmt_time(mean_mk) + util_str = f"{util_pct:.1f}%" + row = (f" {m_str:>4} {policy:<16} {n_stages:>7} " + f"{mk_str:>12} {util_str:>{cpu_col_w}} {peak:>7.0f}MB") + if scalable_names and worker_assignment: + wvals = " ".join( + f"{worker_assignment.get(n, '-'):>12}" + for n in scalable_names + ) + row += f" {wvals}" + print(row) + print() + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + +def optimize_workers( + workflow, + policy_name: str, + cpu_limit: float, + mem_limit: float, + amdahl_models: Dict[str, AmdahlModel], + cpu_fallback_factor: float = 10.0, + task_overhead: float = 0.1, + n_eval_samples: int = 3, + rng_seed: int = 0, + backfill_model: str = "off", + n_backfill: int = 1, + backfill_cpu_factor: float = 1.5, + backfill_mem_factor: float = 1.5, + backfill_slowdown_factor: float = 1.15, + maxjobs: int = 10_000, +) -> Tuple[Dict[str, int], float]: + """Coordinate-descent search for the best worker assignment. + + For each scalable task, iterates over its valid worker range (from the + Amdahl model) and picks the count that minimises mean makespan while + holding all other tasks fixed. Repeats until no improvement is found. + + Returns (best_assignment, best_makespan_s). + """ + # Start from the current assignment implied by each model's cpu_mean_ref + # (what update_resource_estimates already set via round(cpu_mean)). + assignment: Dict[str, int] = { + name: max(model.min_workers, + min(model.max_workers, max(1, round(model.cpu_mean_ref)))) + for name, model in amdahl_models.items() + } + + walltime_stds: Dict[str, float] = {} # no learned std for optimizer runs + + def _evaluate(asgn: Dict[str, int]) -> float: + makespans = [] + for s in range(n_eval_samples): + rng = random.Random(rng_seed + s) if n_eval_samples > 1 else None + r = simulate( + workflow, policy_name, cpu_limit, mem_limit, + cpu_fallback_factor=cpu_fallback_factor, + task_overhead=task_overhead, + walltime_stds=walltime_stds, + cv_fallback=0.1, + rng=rng, + amdahl_models=amdahl_models, + worker_assignment=asgn, + backfill_model=backfill_model, + n_backfill=n_backfill, + backfill_cpu_factor=backfill_cpu_factor, + backfill_mem_factor=backfill_mem_factor, + backfill_slowdown_factor=backfill_slowdown_factor, + maxjobs=maxjobs, + ) + makespans.append(r.makespan) + return statistics.mean(makespans) + + best_score = _evaluate(assignment) + improved = True + while improved: + improved = False + for name, model in amdahl_models.items(): + best_n = assignment[name] + for n in model.worker_range: + if n == best_n: + continue + trial = dict(assignment) + trial[name] = n + score = _evaluate(trial) + if score < best_score - 0.1: # 0.1 s improvement threshold + best_score = score + best_n = n + improved = True + assignment[name] = best_n + return assignment, best_score + + +def build_parser() -> argparse.ArgumentParser: + p = argparse.ArgumentParser( + description="Simulate O2DPG workflow scheduling without running tasks.", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + p.add_argument("-f", "--workflowfile", required=True) + p.add_argument("--update-resources", dest="update_resources", default=None, + metavar="JSON[:FIELDS]", + help="Apply learned resources from JSON (same file as " + "--update-resources in the runner). Enables walltime-based " + "critical path. Optionally restrict which dimensions are " + "patched with a colon-separated field list chosen from " + "{cpu,mem,lifetime}. Example: learned.json:lifetime applies " + "only walltimes, keeping cpu/mem from workflow.json.") + p.add_argument("--cpu-limit", type=float, default=8.0) + p.add_argument("--mem-limit", type=float, default=60000.0, help="in MB") + p.add_argument("--policies", nargs="+", + default=["timeframe", "critical-path", "best-fit"], + choices=["timeframe", "critical-path", "best-fit"]) + p.add_argument("-tt", "--target-tasks", nargs="+", default=["*"]) + p.add_argument("--target-labels", nargs="+", default=[]) + p.add_argument("--walltime-per-core", type=float, default=10.0, metavar="S", + help="Fallback walltime per CPU core [s] when no learned " + "walltime is available.") + p.add_argument("--task-overhead", type=float, default=0.1, metavar="S", + help="Per-task idle overhead [s] added to every task's slot " + "duration (models alienv load, process startup, I/O flush, " + "and scheduler reaction time). This time contributes zero " + "CPU to the utilisation numerator, so larger values lower " + "both predicted makespan and CPU efficiency. " + "Default 0.1 s is conservative; calibration against " + "measurements typically gives 5–7 s for production " + "ALICE MC workflows.") + p.add_argument("--backfill-model", default="off", + choices=["off", "structural", "slowdown", "holefill"], + help="Backfill approximation used by the simulator. " + "'structural' replays the runner's second admission lane; " + "'slowdown' adds a fitted walltime penalty to backfill tasks; " + "'holefill' lets backfill tasks consume only the CPU left idle " + "by foreground tasks, slowing them proportionally.") + p.add_argument("--n-backfill", type=int, default=1, metavar="N", + help="Maximum concurrent backfill tasks when backfill simulation is enabled.") + p.add_argument("--backfill-cpu-factor", type=float, default=1.5, metavar="X", + help="Total CPU oversubscription factor allowed for backfill admission.") + p.add_argument("--backfill-mem-factor", type=float, default=1.5, metavar="X", + help="Total memory oversubscription factor allowed for backfill admission.") + p.add_argument("--backfill-slowdown-factor", type=float, default=1.15, metavar="X", + help="Walltime multiplier applied to backfill tasks in " + "--backfill-model slowdown.") + p.add_argument("--samples", type=int, default=1, metavar="N", + help="Number of Monte Carlo samples for stochastic simulation. " + "When >1, walltime for each task is drawn from a log-normal " + "distribution parameterised by lifetime.mean and lifetime.std " + "from the learned JSON. Output shows mean ± std of makespan.") + p.add_argument("--cv-fallback", type=float, default=0.15, metavar="CV", + help="Coefficient of variation (std/mean) used as a noise floor " + "for tasks whose learned walltime std is zero or absent " + "(single-TF runs, old learned files, etc.). " + "0 disables fallback and makes those tasks deterministic.") + p.add_argument("-j", "--maxjobs", type=int, default=0, metavar="N", + help="Maximum concurrent tasks. -1 = serial mode (1 task at a time): " + "reproduces the learning-run conditions (-jmax 1) and makes the " + "optimizer prefer maximum workers per task. 0 = unlimited (default). " + "N > 0 = at most N concurrent tasks.") + p.add_argument("--optimize-workers", action="store_true", + help="Run coordinate-descent optimizer to find the best worker " + "assignment for scalable tasks. Requires --update-resources " + "with a learned.json that contains 'amdahl' blocks (produced " + "by json-stat --workflow workflow.json).") + p.add_argument("--opt-eval-samples", type=int, default=3, metavar="N", + help="Simulator evaluations per candidate during optimization " + "(more = less noise, slower).") + p.add_argument("--write-optimized", default=None, metavar="FILE", + help="After --optimize-workers, write a copy of the learned JSON " + "with updated lifetime.mean and cpu.mean for scalable tasks. " + "Pass this file to --update-resources in the runner to apply " + "optimized worker counts. When multiple --policies are given " + "the policy with the best (lowest) optimized makespan is used.") + p.add_argument("--verbose", action="store_true", + help="Print per-task schedule for each policy.") + p.add_argument("--output", default=None, metavar="FILE", + help="Write results as JSON to FILE.") + p.add_argument("--timeframes", type=int, nargs="+", default=None, metavar="M", + help="Simulate with M timeframes instead of the workflow's original count. " + "Detects the per-TF template structure automatically and replicates it. " + "Pass multiple values for a sweep (e.g. --timeframes 1 2 5 10 20) to " + "produce a table of makespan and CPU utilisation vs timeframe count.") + return p + + +def main(argv=None) -> int: + ns = build_parser().parse_args(argv) + + raw = load_json(ns.workflowfile) + target_tasks = [t.strip('"').strip("'") for t in ns.target_tasks] + + # ── Parse --update-resources path[:field,field,...] ─────────────────────── + ur_path: Optional[str] = None + ur_fields: Optional[Set[str]] = None # None = all three fields + if ns.update_resources: + _parts = ns.update_resources.split(":", 1) + ur_path = _parts[0] + if len(_parts) > 1: + _raw_fields = {f.strip().lower() for f in _parts[1].split(",")} + _bad = _raw_fields - _LEARN_ALL + if _bad: + print(f"ERROR: unknown field(s) in --update-resources specifier: " + f"{sorted(_bad)}. Valid: cpu, mem, lifetime", file=sys.stderr) + return 1 + ur_fields = _raw_fields + + # ── Load learned JSON once (independent of timeframe count) ────────────── + walltime_stds: Dict[str, float] = {} + learned_full: Dict = {} + amdahl_models: Dict[str, AmdahlModel] = {} + + if ur_path: + _active = ur_fields if ur_fields is not None else _LEARN_ALL + _fields_str = ", ".join(sorted(_active)) + print(f"Applying learned resources from {ur_path}" + + (f" (fields: {_fields_str})" if ur_fields is not None else "") + + " ...") + with open(ur_path) as fh: + learned_full = json.load(fh) + apply_lifetime = "lifetime" in _active + for name, data in learned_full.items(): + if name == "count": + continue + if ns.samples > 1 and apply_lifetime: + std = data.get("lifetime", {}).get("std", 0.0) or 0.0 + if std > 0: + walltime_stds[name] = float(std) + if isinstance(data, dict) and "amdahl" in data: + try: + amdahl_models[name] = AmdahlModel.from_dict(data["amdahl"]) + except (KeyError, ValueError): + pass + if amdahl_models: + print(f" Amdahl models loaded for {len(amdahl_models)} scalable task(s): " + f"{', '.join(sorted(amdahl_models))}") + else: + print(f"No learned resources; using cpu * {ns.walltime_per_core}s as walltime proxy.") + + stochastic = ns.samples > 1 + if stochastic: + n_with_std = len(walltime_stds) + print(f" Stochastic: {n_with_std} tasks use learned std, " + f"cv_fallback={ns.cv_fallback:.2f} for the rest.") + if ns.backfill_model != "off": + print( + "Backfill simulation: " + f"model={ns.backfill_model}, n_backfill={ns.n_backfill}, " + f"cpu_factor={ns.backfill_cpu_factor}, mem_factor={ns.backfill_mem_factor}, " + + ( + f"slowdown={ns.backfill_slowdown_factor:.2f}x\n" + if ns.backfill_model == "slowdown" + else "foreground-hole driven\n" + if ns.backfill_model == "holefill" + else "\n" + ) + ) + + # Verbose Amdahl model summary (shown once, outside the sweep loop). + if amdahl_models and ns.verbose: + print() + for name, model in sorted(amdahl_models.items()): + n_cur = max(model.min_workers, + min(model.max_workers, max(1, round(model.cpu_mean_ref)))) + print(f" Amdahl: {name} " + f"(n_ref={model.n_ref}, cpu_mean={model.cpu_mean_ref:.2f}, " + f"t_serial={model.t_serial:.1f}s, " + f"t_parallel_tot={model.t_parallel_tot:.1f}s)") + print(f" {'n':>4} {'walltime':>10} {'Δ vs current':>14}") + wt_cur = model.walltime(n_cur) + for n in model.worker_range: + wt = model.walltime(n) + marker = " ← current" if n == n_cur else "" + print(f" {n:>4} {_fmt_time(wt):>10} " + f"{wt - wt_cur:>+12.1f}s{marker}") + print() + + # Resolve maxjobs — must come before optimizer and policy loops. + # -j -1 → serial + n_ref workers -j 1 → serial + round(cpu_mean) workers + # -j 0 → unlimited (default) -j N → at most N concurrent tasks + serial_mode = ns.maxjobs == -1 + procs_limit = 1 if (serial_mode or ns.maxjobs == 1) else (10_000 if ns.maxjobs <= 0 else ns.maxjobs) + if serial_mode: + print("Serial mode (-j -1): reproducing learning-run conditions " + "(1 task at a time, n_ref workers per scalable task).") + + # Default worker assignment: Amdahl-based, same for all M values in a sweep. + default_worker_assignment: Optional[Dict[str, int]] = None + if amdahl_models: + if serial_mode: + default_worker_assignment = {name: m.n_ref for name, m in amdahl_models.items()} + parts = [f"{n}: {w}w (n_ref)" for n, w in sorted(default_worker_assignment.items())] + else: + default_worker_assignment = { + name: max(m.min_workers, min(m.max_workers, max(1, round(m.cpu_mean_ref)))) + for name, m in amdahl_models.items() + } + parts = [f"{n}: {w}w" for n, w in sorted(default_worker_assignment.items())] + print(f"Worker counts for simulation: {', '.join(parts)}\n") + + # ── Timeframe sweep ─────────────────────────────────────────────────────── + tf_list: List[Optional[int]] = sorted(set(ns.timeframes)) if ns.timeframes else [None] + sweep_mode = len(tf_list) > 1 + if sweep_mode: + print(f"Timeframe sweep: M={tf_list} policies={ns.policies} " + f"cpu_limit={ns.cpu_limit} mem_limit={ns.mem_limit} MB\n") + + all_sweep_results: List[Tuple] = [] # (M, results_by_policy, n_stages, worker_assignment) + last_opt_results: List[Tuple[str, Dict[str, int], float]] = [] + + for M in tf_list: + if sweep_mode: + print(f" M={M} ...", end=" ", flush=True) + + cur_raw = replicate_workflow_for_timeframes(raw, M) if M is not None else raw + wf = build_workflow(cur_raw, target_tasks, ns.target_labels) + if not wf.stages: + print(f"Workflow is empty after filtering (M={M}).") + continue + + if ur_path: + if ur_fields is None: + update_resource_estimates(wf, ur_path) + else: + _apply_learned_fields(wf, learned_full, ur_fields) + has_wt = sum(1 for t in wf.stages if t.get("resources", {}).get("walltime")) + if not sweep_mode: + print(f" {has_wt}/{len(wf.stages)} tasks have learned walltime.") + if stochastic: + n_fallback = len(wf.stages) - len(walltime_stds) + print(f" Stochastic: {len(walltime_stds)} tasks use learned std, " + f"{n_fallback} use cv_fallback={ns.cv_fallback:.2f}.") + if not sweep_mode: + print(f"Workflow: {len(wf.stages)} tasks, cpu_limit={ns.cpu_limit}, " + f"mem_limit={ns.mem_limit} MB, samples={ns.samples}" + + (" (stochastic)" if stochastic else " (deterministic)") + "\n") + + # Worker-count optimizer (--optimize-workers): runs per-M so that the + # simulated workflow matches the actual task count. + cur_worker_assignment = default_worker_assignment + opt_results: List[Tuple[str, Dict[str, int], float]] = [] + if ns.optimize_workers: + if not amdahl_models: + print("WARNING: --optimize-workers requires Amdahl models in learned.json. " + "Re-run json-stat with --workflow workflow.json first.", file=sys.stderr) + else: + if not sweep_mode: + print(f"\nOptimizing worker assignment over " + f"{len(amdahl_models)} scalable task(s)...") + for policy_name in ns.policies: + best_asgn, best_mk = optimize_workers( + wf, policy_name, ns.cpu_limit, ns.mem_limit, + amdahl_models=amdahl_models, + cpu_fallback_factor=ns.walltime_per_core, + task_overhead=ns.task_overhead, + n_eval_samples=ns.opt_eval_samples, + backfill_model=ns.backfill_model, + n_backfill=ns.n_backfill, + backfill_cpu_factor=ns.backfill_cpu_factor, + backfill_mem_factor=ns.backfill_mem_factor, + backfill_slowdown_factor=ns.backfill_slowdown_factor, + maxjobs=procs_limit, + ) + opt_results.append((policy_name, best_asgn, best_mk)) + if not sweep_mode: + print(f"\n [{policy_name}] best makespan: {_fmt_time(best_mk)}") + print(f" {'Task':<35} {'n_ref':>6} {'current':>8} {'→ opt':>6} " + f"{'wt(current)':>12} {'wt(opt)':>10} {'Δwt':>8}") + print(f" {'-'*85}") + for name, model in sorted(amdahl_models.items()): + n_cur = max(model.min_workers, + min(model.max_workers, max(1, round(model.cpu_mean_ref)))) + n_opt = best_asgn[name] + wt_cur = model.walltime(n_cur) + wt_opt = model.walltime(n_opt) + delta = wt_opt - wt_cur + changed = "←" if n_opt != n_cur else "" + print(f" {name:<35} {model.n_ref:>6} {n_cur:>8} {n_opt:>6} " + f"{_fmt_time(wt_cur):>12} {_fmt_time(wt_opt):>10} " + f"{delta:>+7.1f}s {changed}") + if not sweep_mode: + print() + _, best_asgn_opt, _ = min(opt_results, key=lambda x: x[2]) + cur_worker_assignment = best_asgn_opt + last_opt_results = opt_results + + # ── Run simulations ─────────────────────────────────────────────────── + results_by_policy: Dict[str, List[SimResult]] = {} + for policy_name in ns.policies: + runs: List[SimResult] = [] + for s in range(ns.samples): + rng = random.Random(s) if stochastic else None + r = simulate( + wf, policy_name, + cpu_limit=ns.cpu_limit, + mem_limit=ns.mem_limit, + cpu_fallback_factor=ns.walltime_per_core, + task_overhead=ns.task_overhead, + walltime_stds=walltime_stds, + cv_fallback=ns.cv_fallback, + rng=rng, + amdahl_models=amdahl_models if amdahl_models else None, + worker_assignment=cur_worker_assignment, + backfill_model=ns.backfill_model, + n_backfill=ns.n_backfill, + backfill_cpu_factor=ns.backfill_cpu_factor, + backfill_mem_factor=ns.backfill_mem_factor, + backfill_slowdown_factor=ns.backfill_slowdown_factor, + maxjobs=procs_limit, + ) + runs.append(r) + results_by_policy[policy_name] = runs + if not sweep_mode: + makespans = [r.makespan for r in runs] + summary = _fmt_time(statistics.mean(makespans)) + if stochastic: + summary += f" ± {_fmt_time(statistics.stdev(makespans))}" + print(f" {policy_name:<16} makespan={summary}") + + all_sweep_results.append((M, results_by_policy, len(wf.stages), cur_worker_assignment)) + + if sweep_mode: + best_mk = min( + statistics.mean(r.makespan for r in runs) + for runs in results_by_policy.values() + ) + print(f"{len(wf.stages)} tasks, best makespan={_fmt_time(best_mk)}") + + if not all_sweep_results: + return 1 + + # ── Summary ─────────────────────────────────────────────────────────────── + if sweep_mode: + _print_sweep_table(all_sweep_results, ns.cpu_limit, ns.samples) + else: + _, results_by_policy, _, _ = all_sweep_results[0] + print_summary(results_by_policy, ns.cpu_limit, ns.samples) + if ns.verbose: + for policy_name, runs in results_by_policy.items(): + print_verbose(runs[0]) + + # ── Write optimized learned.json ────────────────────────────────────────── + if ns.write_optimized and last_opt_results and learned_full: + if sweep_mode: + print(f"NOTE: --write-optimized in sweep mode writes the optimisation " + f"result from the last M={tf_list[-1]}.") + best_policy, best_asgn, best_mk = min(last_opt_results, key=lambda x: x[2]) + if len(ns.policies) > 1: + print(f"Writing optimized resources: policy '{best_policy}' " + f"selected (best makespan {_fmt_time(best_mk)}).") + else: + print(f"Writing optimized resources (makespan {_fmt_time(best_mk)}).") + out_learned = copy.deepcopy(learned_full) + n_changed = 0 + for name, model in amdahl_models.items(): + if name not in out_learned: + continue + n_opt = best_asgn[name] + n_cur = max(model.min_workers, + min(model.max_workers, max(1, round(model.cpu_mean_ref)))) + if "lifetime" not in out_learned[name]: + out_learned[name]["lifetime"] = {} + out_learned[name]["lifetime"]["mean"] = round(model.walltime(n_opt), 3) + if "cpu" not in out_learned[name]: + out_learned[name]["cpu"] = {} + out_learned[name]["cpu"]["mean"] = float(n_opt) + if n_opt != n_cur: + n_changed += 1 + with open(ns.write_optimized, "w") as fh: + json.dump(out_learned, fh, indent=2) + print(f" Wrote {ns.write_optimized} " + f"({len(amdahl_models)} scalable tasks updated, " + f"{n_changed} changed from current assignment).") + + # ── JSON output ─────────────────────────────────────────────────────────── + if ns.output: + out_list = [] + for M, results_by_policy, n_stages, _ in all_sweep_results: + for policy_name, runs in results_by_policy.items(): + for i, r in enumerate(runs): + d = r.to_dict() + d["sample"] = i + d["n_stages"] = n_stages + if M is not None: + d["timeframes"] = M + out_list.append(d) + with open(ns.output, "w") as fh: + json.dump(out_list, fh, indent=2) + print(f"Results written to {ns.output}") + + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/MC/workflow_runner/o2dpg_workflow_runner.py b/MC/workflow_runner/o2dpg_workflow_runner.py new file mode 100755 index 000000000..70c56a4ba --- /dev/null +++ b/MC/workflow_runner/o2dpg_workflow_runner.py @@ -0,0 +1,18 @@ +#!/usr/bin/env python3 +"""Entry point preserving the legacy filename MC/bin/o2dpg_workflow_runner.py. + +All behavior lives in the o2dpg_runner package next to this script. +""" + +import os +import sys + +# Make sure we can import the o2dpg_runner package that sits next to us. +_here = os.path.dirname(os.path.realpath(__file__)) +if _here not in sys.path: + sys.path.insert(0, _here) + +from o2dpg_runner.cli import main + +if __name__ == "__main__": + sys.exit(main())