Skip to content

Adam #2

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
rajarshitiwari wants to merge 2 commits into
base: master
Choose a base branch
from
Open

Adam #2

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
b9973f0
Replaced with Adam's version tested on Meluxina
rajarshitiwari Jul 30, 2026
24757ea
taskfarm: 3.1 version from Adam
rajarshitiwari Jul 30, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
174 changes: 125 additions & 49 deletions taskfarm
100755 → 100644
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
#!/usr/bin/env python3
#!/apps/USE/easybuild/release/2023.1/software/Python/3.11.3-GCCcore-12.3.0/bin/python -E
#
# ICHEC Taskfarm utility
#
Expand All @@ -15,6 +15,8 @@
# Fionn inherits taskfarm2/1.2 from Stokes, with some minor modifications. But
# the name of the module reverts to taskfarm with an initial version number of
# 2.3 to reflect its provenance.
#
# v2.8 -> v2.9 marks the transition from Fionn (PBS) to Kay (Slurm)
#
# Version 1.0: Initial release on Stokes.
#
Expand All @@ -41,9 +43,9 @@
# task should inherit its environment from the PBS script calling
# taskfarm; sourcing startup files might cause unexpected behaviour.
# Version 2.4: New features
# - force a kill of the farm if a certain file exists in the forder of the running task
# - force a kill of the farm if a certain file exists in the folder of the running task
# default name of the file is abbadon change it via TASKFARM_STOPFILE
# - force a kill of the farm if a certain file exists in the forder of the running task
# - force a kill of the farm if a certain file exists in the folder of the running task
# and contains a certain magic word by default no magic is set.
# use TASKFARM_STOPMAGIC to define a word
# - TASKFARM_SLEEP controls at what intervals one checks for the file and, if the case,
Expand All @@ -63,58 +65,110 @@
# - Remove empty lines from the task list.
# - Use I_MPI_PIN_PROCESSOR_LIST instead of numactl, in order to avoid clash between
# cores selected by the mpirun and the one required via numactl.
# Version 2.8: Bug fix
# - The I_MPI_PIN_PROCESSOR_LIST does not allow multi-threading, since it pins only
# one core per MPI rank. To fix this, in case of not-MPI task, replacing the list
# with a numactl and I_MPI_PIN_DOMAIN combo.
# I_MPI_PIN_DOMAIN=node to allow all cores, then numactl with same list as
# I_MPI_PIN_PROCESSOR_LIST one
# Version 2.9: Initial port from Fionn to Kay
# - Changed the launcher from mpirun to mpiexec ("mpiexec -bootstrap slurm" is same
# as mpirun, which sets other variables aimed at MPI jobs that are not needed here
# and *may* allow users launch multiple MPI tasks in a cleaner environment).
# - Changed $PBS_JOBID to $SLURM_JOBID
# - Parse $SLURM_JOB_NODELIST env var and re-factored to simulate prior $PBS_NODEFILE
# Version 2.9: Bug fix without new release.
# - Correctly calculate the number of threads per core.
# Version 2.9: Bug fix without new release (15/03/2019)
# - Fixed bug: exp_nodelist has one entry per node (SLURM) rather than multiple
# entries (one per core) per node as in $PBS_NODEFILE. Bug was introduced when
# migrating from Fionn to Kay.
# Version 2.9: Bug fix without new release (17/11/2021) AR
# - SLURM was updated on Kay. Interactivate runs are now generated by `salloc` and not
# `srun`. In using the -bootstrap option `srun` is called by default, the consequence
# of this is that only a single task can run at any one time. It is possible to change
# with -bootstrap-exec and -bootstrap-exec-args. However it appears that not using
# -bootstrap option seems to work, in the "launch" variable.
#
#
# Version 3.0: changed script to work on Meluxina (AR)
# Changed to work with python 3, e.g. altered print statements.
# Removed the adding of envirinment variables to task scripts.
# There was a syntax error that prevented them running.
# Had to convert some variables to int from float.
# Only works on 1 node, perhaps restrictions on number of ssh connections
# Version 3.1: major changes to make work with srun (AR)
# Added TASKFARM_PROCS_PER_TASK for multi-threaded apps
# Added TASKFARM_MPI_PROCS for MPI enabled apps
# SLURM vars set from TASKFARM environment variables
# Now works on multiple nodes

import os
import sys
import stat
import signal
import time
import re

from subprocess import Popen
from subprocess import Popen, PIPE

# Physical characteristics of the compute nodes
lines=[ line.strip().replace('\t','').split(':') for line in open('/proc/cpuinfo') ]
threads_per_core = len(set([ line[1] for line in lines if line[0]=='physical id']))
siblings = set([ line[1] for line in lines if line[0]=='siblings'])
cores = set([ line[1] for line in lines if line[0]=='cpu cores'])
threads_per_core = int(siblings.pop()) / int(cores.pop())
cores_per_node = len(set([ line[1] for line in lines if line[0]=='processor']))/threads_per_core
# Result of above line is a float, so change to int.
cores_per_node = int(cores_per_node)


# A print function to control verbosity
def verbose_print(string):
def verbose_print(str):
if not 'TASKFARM_SILENT' in os.environ:
print(string)
print(str)
sys.stdout.flush()

# Check comand line arguments
if len(sys.argv) != 2:
sys.stderr.write('Usage: %s <tasklist>' % sys.argv[0] + '\n')
if len(sys.argv) != 2 or sys.argv[1]=="-h":
print('Usage: taskfarm <tasklist>',file=sys.stderr,flush=True)
print('TASKFARM_PPN: number of tasks per node (def=cores per node)',file=sys.stderr,flush=True)
print('TASKFARM_PROCS_PER_TASK: number of threads per task (def=1)',file=sys.stderr,flush=True)
print('TASKFARM_MPI_PROCS: number of MPI procs per task (def=1)',file=sys.stderr,flush=True)
sys.exit(1)
taskfile = sys.argv[1]

# How many processes per node
if 'TASKFARM_SMT' in os.environ:
smt = threads_per_core
else:
smt = 1

if 'TASKFARM_PPN' in os.environ:
try:
ppn = int(os.environ['TASKFARM_PPN'])
except:
sys.stderr.write('Error: $TASKFARM_PPN must be an integer value.\n')
print('Error: $TASKFARM_PPN must be an integer value.a',file=sys.stderr,flush=True)
sys.exit(1)
else:
ppn = cores_per_node


if 'TASKFARM_PROCS_PER_TASK' in os.environ:
try:
cores_per_task = int(os.environ['TASKFARM_PROCS_PER_TASK'])
except:
print('Error: $TASKFARM_RPOCS_PER_TASK must be an integer value.a',file=sys.stderr,flush=True)
sys.exit(1)
else:
ppn = cores_per_node * smt
cores_per_task = 1

if 'TASKFARM_MPI' in os.environ:
cores_per_task=cores_per_node/ppn
if 'TASKFARM_MPI_PROCS' in os.environ:
try:
mpi_per_task = int(os.environ['TASKFARM_MPI_PROCS'])
except:
print('Error: $TASKFARM_MPI_PROCS must be an integer value',file=sys.stderr,flush=True)
sys.exit(1)
else:
cores_per_task=1
mpi_per_task = 1


if 'TASKFARM_MPI_LAUNCHER' in os.environ:
launch=os.environ['TASKFARM_MPI_LAUNCHER']
else:
launch='mpirun'
launch='srun'

if 'TASKFARM_SLEEP' in os.environ:
sleep = float(os.environ['TASKFARM_SLEEP'])
Expand All @@ -141,43 +195,61 @@ else:
keep=False

# Error if an invalid process count is requested
if smt > 1 and ppn > cores_per_node * threads_per_core:
sys.stderr.write('Error: $TASKFARM_PPN must not exceed %d processes per node when $TASKFARM_SMT is set.' %(cores_per_node * threads_per_core,) + '\n')
total_cores_per_node = ppn*cores_per_task*mpi_per_task
verbose_print(' total_cores = %d' % (total_cores_per_node))
if total_cores_per_node > cores_per_node:
print('Error:cores per node must not exceed {0} processes per node.'.format(cores_per_node),file=sys.stderr,flush=True)
sys.exit(1)
elif total_cores_per_node < 1:
print('Error: $TASKFARM_PPN must request one or more processes per node.',file=sys.stderr,flush=True)
sys.exit(1)
elif smt == 1 and ppn > cores_per_node:
sys.stderr.write('Error: $TASKFARM_PPN must not exceed %d processes per node.' %(cores_per_node,) + '\n')


if 'SLURM_JOBID' not in os.environ:
print('Error: no $SLURM_JOBID environment variable found. Do not run taskfarm on login nodes.',file=sys.stderr,fluysh=True)
sys.exit(1)
elif ppn < 1:
sys.stderr.write('Error: $TASKFARM_PPN must request one or more processes per node.\n')
if 'SLURM_JOB_NODELIST' not in os.environ:
print('Error: no $SLURM_JOB_NODELIST environment variable found. Do not run taskfarm on login nodes.',file=sys.stderr,flush=True)
sys.exit(1)


# Generate a list of unique nodes
work_nodes = []
node_ids = []
id_map = {}
core_range={}
jid=os.environ['SLURM_JOBID']
cpath=os.environ['PWD']

try:
nodef = os.environ['SLURM_JOB_NODELIST']
except:
sys.stderr.write('Error opening SLURM_JOB_NODELIST. Exiting.\n')
nodelist_command = 'scontrol show hostnames ' + os.environ['SLURM_JOB_NODELIST']
nodelist_command = nodelist_command.split(' ')
popen_tmp = Popen(nodelist_command, stdout=PIPE,encoding='utf8')
exp_nodelist = popen_tmp.stdout.read().split('\n')
exp_nodelist = filter(None, exp_nodelist)
exp_nodelist = [x for x in exp_nodelist for i in range(cores_per_node)]


except KeyError:
print('Error interpreting the list of nodes from $SLURM_JOB_NODELIST. Exiting.',file=sys.stderr,flush=True)
sys.exit(2)

id = 0
cpt=cores_per_node/ppn
for line in nodef:
node = line.strip()

for node in exp_nodelist:
if work_nodes.count(node) == 0:
for i in range(ppn):
tmp_id = id+i%cores_per_node
node_ids.append(tmp_id)
work_nodes.append(node)
id_map[tmp_id] = node
start=i%cores_per_node*cpt
end=start+cpt-1
core_range[tmp_id]=str(start)+"-"+str(end)
start=int(i%cores_per_node*cpt)
end=int(start+cpt)
lrng = list(range(start,end))
core_range[tmp_id]=','.join(str(scpu) for scpu in lrng)
id += 1
nodef.close()

verbose_print('Taskfarm started with %d workers (%d per node).' % (len(work_nodes), ppn))

def extractPaths(taskline):
Expand Down Expand Up @@ -207,7 +279,7 @@ tasknum = 0
try:
taskf = open(taskfile)
except:
sys.stderr.write('Error opening task file. Exiting.\n')
print('Error opening task file. Exiting.',file=sys.stderr,flush=True)
sys.exit(2)

l1 = [ processLine(line.strip()) for line in taskf if len(line.strip()) > 0]
Expand All @@ -226,16 +298,17 @@ if groupTasks == 1 :
else:
verbose_print('Taskfarm read %d tasks from file \'%s\' grouped as %d metatasks.' % (tf,taskfile,len(tasklist)))

if len(tasklist) % len(work_nodes) != 0:
verbose_print('Warning: Taskfarm input should ideally provide a multiple of %d tasks for %d workers.' % (len(work_nodes), len(work_nodes)))
if len(tasklist) / len(work_nodes) > 20:
verbose_print('Warning: There are %d tasks for %d workers. Taskfarm is not ideal for high-throughput workloads.' % (len(tasklist), len(work_nodes)))
verbose_print('Warning: Running many tasks of a very short duration with Taskfarm is quite inefficient.')
verbose_print('Info: You can aggregate tasks using export TASKFARM_GROUP=xxx')
verbose_print('Info: with xxx how many consecutive tasks to group in a metatask')

# Build an environment
environ = ''
os.environ['SLURM_NTASKS_PER_NODE']=str(ppn*mpi_per_task)
os.environ['SLURM_TASKS_PER_NODE']=str(ppn*mpi_per_task)
os.environ['SLURM_CPUS_PER_TASK']=str(cores_per_task)
os.environ['SLURM_NTASKS']=str(len(tasklist*mpi_per_task))
for param in os.environ:
if param not in ['PROFILEREAD', 'BASH_FUNC_module()'] :
environ = environ + 'export ' + param + '=\'' + os.environ[param] + '\'; '
Expand All @@ -254,9 +327,7 @@ def checkMagic(filepath):
for line in f:
if stopMagic in line:
return True

return False


def checkUserExit(pid):
fpath=EnforceTrailingSlash(taskinfo[pid]['path'])+stopfile
Expand All @@ -280,7 +351,7 @@ def wait():
checkUserExit(pid)

if len(finished) > 0:
pid = finished.keys()[0]
pid = list(finished.keys())[0]
exit = signal = finished[pid]['status']
stillSleepy=False
else:
Expand All @@ -290,7 +361,7 @@ def wait():
# exit = status >> 8
id = taskinfo[pid]['id']
if exit != 0:
sys.stderr.write("'%s' killed by sig %d" % (taskinfo[pid]['task'], signal) + '\n')
print("{0} killed by sig {1}".format(taskinfo[pid]['task'], signal),file=sys.stderr,flush=True)
if not keep:
popen_tmp = os.unlink(taskinfo[pid]['script'])
del taskinfo[pid]
Expand All @@ -311,18 +382,23 @@ for task in tasklist:
host = id_map[id]
cores=core_range[id]
del node_ids[0]
task = environ + ' cd ' + os.environ['PWD'] + ' && ' + task
# task = environ + ' cd ' + os.environ['PWD'] + ' && ' + task
task = ' cd ' + os.environ['PWD'] + ' && ' + task
fp=cpath+"/task-"+host+"-id"+str(id)+"-"+jid + '.' + str(k)
f=open(fp,"w")
f.write("#!/bin/bash\n");
f.write(task)
f.close()
os.chmod(fp,0o755)
command = "%s -n %d -env I_MPI_PIN_PROCESSOR_LIST %s -host %s %s" % (launch,cores_per_task,cores,host,fp)
popen_tmp = Popen([launch,'-env','I_MPI_PIN_PROCESSOR_LIST',cores,'-n', str(cores_per_task),'-host',host,fp])
os.chmod(fp,stat.S_IRWXU)
command = "%s -N 1 -n %d -c %d --nodelist=%s --exact %s" % (launch,mpi_per_task,cores_per_task,host,fp)
verbose_print(command)

command = command.split(' ')
popen_tmp = Popen(command)
pid = popen_tmp.pid
taskinfo[pid] = {'id': id, 'task': command, 'process': popen_tmp, 'path':cpath+"/"+extraPaths[k],'script':fp}
k += 1

# Once all tasks have been started wait for them to finish
while len(taskinfo) > 0:
wait()
Expand Down