diff --git a/lib/python/ray/__init__.py b/lib/python/ray/__init__.py index 075d41541..629908dbd 100644 --- a/lib/python/ray/__init__.py +++ b/lib/python/ray/__init__.py @@ -13,6 +13,7 @@ if hasattr(ctypes, "windll"): # This is done by associating all child processes with a "job" object that imposes this behavior. (lambda kernel32: (lambda job: (lambda n: kernel32.SetInformationJobObject(job, 9, "\0" * 17 + chr(0x8 | 0x4 | 0x20) + "\0" * (n - 18), n))(0x90 if ctypes.sizeof(ctypes.c_void_p) > ctypes.sizeof(ctypes.c_int) else 0x70) and kernel32.AssignProcessToJobObject(job, ctypes.c_void_p(kernel32.GetCurrentProcess())))(ctypes.c_void_p(kernel32.CreateJobObjectW(None, None))) if kernel32 is not None else None)(ctypes.windll.kernel32) +import ray.experimental import ray.serialization from ray.worker import register_class, error_info, init, connect, disconnect, get, put, wait, remote from ray.worker import Reusable, reusables diff --git a/lib/python/ray/experimental/__init__.py b/lib/python/ray/experimental/__init__.py new file mode 100644 index 000000000..e515c8dc5 --- /dev/null +++ b/lib/python/ray/experimental/__init__.py @@ -0,0 +1,5 @@ +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from .utils import copy_directory diff --git a/lib/python/ray/experimental/utils.py b/lib/python/ray/experimental/utils.py new file mode 100644 index 000000000..7c8083214 --- /dev/null +++ b/lib/python/ray/experimental/utils.py @@ -0,0 +1,74 @@ +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import io +import os +import tarfile +import sys + +import ray + +def tarred_directory_as_bytes(source_dir): + """Tar a directory and return it as a byte string. + + Args: + source_dir (str): The name of the directory to tar. + + Returns: + A byte string representing the tarred file. + """ + # Get a BytesIO object. + string_file = io.BytesIO() + # Create an in-memory tarfile of the source directory. + with tarfile.open(mode="w:gz", fileobj=string_file) as tar: + tar.add(source_dir, arcname=os.path.basename(source_dir)) + string_file.seek(0) + return string_file.read() + +def tarred_bytes_to_directory(tarred_bytes, target_dir): + """Take a byte string and untar it. + + Args: + tarred_bytes (str): A byte string representing the tarred file. This should + be the output of tarred_directory_as_bytes. + target_dir (str): The directory to create the untarred files in. + """ + string_file = io.BytesIO(tarred_bytes) + with tarfile.open(fileobj=string_file) as tar: + tar.extractall(path=target_dir) + +def copy_directory(source_dir, target_dir=None): + """Copy a local directory to each machine in the Ray cluster. + + Note that both source_dir and target_dir must have the same basename). For + example, source_dir can be /a/b/c and target_dir can be /d/e/c. In this case, + the directory /d/e will be added to the Python path of each worker. + + Note that this method is not completely safe to use. For example, workers that + do not do the copying and only set their paths (only one worker per node does + the copying) may try to execute functions that use the files in the directory + being copied before the directory being copied has finished untarring. + + Args: + source_dir (str): The directory to copy. + target_dir (str): The location to copy it to on the other machines. If this + is not provided, the source_dir will be used. If it is provided and is + different from source_dir, the source_dir also be copied to the target_dir + location on this machine. + """ + target_dir = source_dir if target_dir is None else target_dir + source_dir = os.path.abspath(source_dir) + target_dir = os.path.abspath(target_dir) + source_basename = os.path.basename(source_dir) + target_basename = os.path.basename(target_dir) + if source_basename != target_basename: + raise Exception("The source_dir and target_dir must have the same base name, {} != {}".format(source_basename, target_basename)) + tarred_bytes = tarred_directory_as_bytes(source_dir) + def f(worker_info): + if worker_info["counter"] == 0: + tarred_bytes_to_directory(tarred_bytes, os.path.dirname(target_dir)) + sys.path.append(os.path.dirname(target_dir)) + # Run this function on all workers to copy the directory to all nodes and to + # add the directory to the Python path of each worker. + ray.worker.global_worker.run_function_on_all_workers(f) diff --git a/test/runtest.py b/test/runtest.py index 6c9a5f5dc..bf69c02c5 100644 --- a/test/runtest.py +++ b/test/runtest.py @@ -2,10 +2,12 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import os import unittest import ray import numpy as np import time +import shutil import string import sys from collections import namedtuple @@ -673,5 +675,53 @@ class ReusablesTest(unittest.TestCase): ray.worker.cleanup() +class UtilsTest(unittest.TestCase): + + def testCopyingDirectory(self): + # The functionality being tested here is really multi-node functionality, + # but this test just uses a single node. + + ray.init(start_ray_local=True, num_workers=1) + + source_text = "hello world" + + temp_dir1 = os.path.join(os.path.dirname(__file__), "temp_dir1") + source_dir = os.path.join(temp_dir1, "dir") + source_file = os.path.join(source_dir, "file.txt") + temp_dir2 = os.path.join(os.path.dirname(__file__), "temp_dir2") + target_dir = os.path.join(temp_dir2, "dir") + target_file = os.path.join(target_dir, "file.txt") + + def remove_temporary_files(): + if os.path.exists(temp_dir1): + shutil.rmtree(temp_dir1) + if os.path.exists(temp_dir2): + shutil.rmtree(temp_dir2) + + # Remove the relevant files if they are left over from a previous run of + # this test. + remove_temporary_files() + + # Create the source files. + os.mkdir(temp_dir1) + os.mkdir(source_dir) + with open(source_file, "w") as f: + f.write(source_text) + + # Copy the source directory to the target directory. + ray.experimental.copy_directory(source_dir, target_dir) + time.sleep(0.5) + + # Check that the target files exist and are the same as the source files. + self.assertTrue(os.path.exists(target_dir)) + self.assertTrue(os.path.exists(target_file)) + with open(target_file, "r") as f: + self.assertEqual(f.read(), source_text) + + # Remove the relevant files to clean up. + remove_temporary_files() + + ray.worker.cleanup() + if __name__ == "__main__": unittest.main(verbosity=2)