diff --git a/worker.py b/worker.py index 546b6fc..1eb4f97 100644 --- a/worker.py +++ b/worker.py @@ -9,6 +9,7 @@ import pynvml import csv import queue +import sys from runtime.rpc import scheduler_server from task import Task, JobInfo @@ -238,7 +239,10 @@ def close(self): parser.add_argument('--log_path', type=str, required=True) # default='results/test_tgs_results.csv') parser.add_argument('--need_throughput', action='store_true', default=False) args = parser.parse_args() - + p = input('Containers will be stopped and maybe deleted. Press y and then enter: ') + if p != 'y': + print('Abort.') + sys.exit(-1) subprocess.call('docker stop $(docker ps -q)', shell=True) subprocess.call('docker rm $(docker ps -aq)', shell=True) @@ -307,4 +311,4 @@ def close(self): sleep_time = min(sleep_time, (worker._start_time + worker._submit_queue[0]['submit_time'] - time.time())) time.sleep(sleep_time) - worker.close() \ No newline at end of file + worker.close()