workers = [
# worker which sync periodically and at the end of task
SyncSSHWorker(
host='localhost', # alias from .ssh/config
sync_period=600, # run rsync every 5 min
remote_root=remote_tmp_root, # where to place task folders on remote
remote_user_rc="export CUDA_VISIBLE_DEVICES=0\n" # enforce use of first GPU,
"export WORKER_ID=SyncSSHWorker-0; \n" # mark task's worker
+ source_venv
),
# worker which sync only at the end of task
SSHWorker(
host='localhost', # alias from .ssh/config
remote_root=remote_tmp_root, # where to place task folders on remote
remote_user_rc="export CUDA_VISIBLE_DEVICES=1\n" # enforce use of second GPU
"export WORKER_ID=SSHWorker-0;\n" # mark task's worker
+ source_venv
),
# simple worker which runs locally, no folder copying, no ssh. Runs as multiprocessing.Process
LocalCudaWorker(
device=3 # use third GPU
)
]
tasks = [
Task.load(task_file)
for task_file in Path("./").glob("*/state.dill") # find all task in subdirectories
]
Pool(workers).run(tasks)
from remote_runner.utility import ChangeToTemporaryDirectory
class MyCommand(Command):
executable = ['python']
def __init__(self, count, line):
super(MyCommand, self).__init__()
self.script = StringArgument(
'-c', f'for i in range({count}): print("{line}")')
class HeavyStdoutOutput(Step):
def run(self, md: 'MdProtocol'):
cmd = MyCommand(64 * 1024, "#" * 256)
cmd.run()
class MyMdProtocol(MdProtocol):
def __init__(self, name: str):
wd = Path(name)
wd.mkdir(mode=0o755, exist_ok=True)
MdProtocol.__init__(self, name, wd)
self.greet = HeavyStdoutOutput("greet")
with ChangeToTemporaryDirectory():
Pool([LocalWorker()]).run([MyMdProtocol("./")])
with open("stdout") as f:
assert f.read() == ""