mirror of
https://github.com/clearml/clearml
synced 2025-01-31 09:07:00 +00:00
96 lines
3.7 KiB
Python
96 lines
3.7 KiB
Python
# ClearML - example of multiple sub-processes interacting and reporting to a single master experiment
|
|
|
|
import multiprocessing
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
from argparse import ArgumentParser
|
|
from random import randint
|
|
|
|
from clearml import Task
|
|
|
|
|
|
def mp_worker(arguments):
|
|
print('sub process', os.getpid())
|
|
data, extra_dict = arguments
|
|
inputs, the_time = data
|
|
Task.current_task().connect(extra_dict)
|
|
print(" Process %s\tWaiting %s seconds" % (inputs, the_time))
|
|
time.sleep(int(the_time))
|
|
print(" Process %s\tDONE" % inputs)
|
|
|
|
|
|
def mp_handler(use_subprocess, data, additional_parameters):
|
|
if use_subprocess:
|
|
process = multiprocessing.Pool(4)
|
|
else:
|
|
process = multiprocessing.pool.ThreadPool(4)
|
|
|
|
process.map(mp_worker, zip(data, additional_parameters))
|
|
process.close()
|
|
print('DONE main !!!')
|
|
|
|
|
|
if __name__ == '__main__':
|
|
parser = ArgumentParser()
|
|
parser.add_argument('--num-workers', help='integer value', type=int, default=3)
|
|
parser.add_argument('--use-subprocess', help="Use sub processes", dest='subprocess', action='store_true')
|
|
parser.add_argument('--no-subprocess', help="Use threads", dest='subprocess', action='store_false')
|
|
parser.add_argument('--additional-parameters', help='task parameters', type=str, nargs="+")
|
|
parser.set_defaults(subprocess=True)
|
|
# this argument we will not be logging, see below Task.init
|
|
parser.add_argument('--counter', help='integer value', type=int, default=-1)
|
|
|
|
args = parser.parse_args()
|
|
print(os.getpid(), 'ARGS:', args)
|
|
|
|
# Fake data for us to "process"
|
|
data = (
|
|
['a', '2'], ['b', '4'], ['c', '6'], ['d', '8'],
|
|
['e', '1'], ['f', '3'], ['g', '5'], ['h', '7'],
|
|
)
|
|
|
|
if not args.additional_parameters:
|
|
# Random task parameters
|
|
args.additional_parameters = [
|
|
f"stuff_{str(randint(0, 100))}:some stuff {str(randint(0, 100))}" for _ in range(len(data))]
|
|
task_parameters = (dict([p.partition(":")[::2]]) for p in args.additional_parameters)
|
|
|
|
# We have to initialize the task in the master process,
|
|
# it will make sure that any sub-process calling Task.init will get the master task object
|
|
# notice that we exclude the `counter` argument, so we can launch multiple sub-processes with clearml-agent
|
|
# otherwise, the `counter` will always be set to the original value.
|
|
task = Task.init('examples', 'Popen example', auto_connect_arg_parser={'counter': False})
|
|
|
|
# we can connect multiple dictionaries, each from different process, as long as the keys have different names
|
|
param = {'args_{}'.format(args.num_workers): 'some value {}'.format(args.num_workers)}
|
|
task.connect(param)
|
|
|
|
# check if we need to start the process, meaning counter is negative
|
|
counter = args.num_workers if args.counter < 0 else args.counter
|
|
|
|
p = None
|
|
# launch sub-process, every subprocess will launch the next in the chain, until we launch them all.
|
|
# We could also launch all of them here, but that would have been to simple for us J
|
|
if counter > 0:
|
|
cmd = [sys.executable, sys.argv[0],
|
|
'--counter', str(counter - 1),
|
|
'--num-workers', str(args.num_workers),
|
|
'--use-subprocess' if args.subprocess else '--no-subprocess',
|
|
'--additional-parameters', *args.additional_parameters]
|
|
|
|
print(cmd)
|
|
p = subprocess.Popen(cmd, cwd=os.getcwd())
|
|
|
|
# the actual "processing" is done here
|
|
mp_handler(args.subprocess, data, task_parameters)
|
|
print('Done logging')
|
|
|
|
# wait for the process we launched
|
|
# this means every subprocess will be waiting for the process it launched and
|
|
# the master process will exit after all of them are completed
|
|
if p and counter > 0:
|
|
p.wait()
|
|
print('Exiting')
|