Repository: mars-project/mars

Issue:
[ENH] Allow specifying plasma_directory
We need to add an argument for Mars worker to pass customized ``plasma_directory`` into ``plasma.start_plasma_store``. What's more, when this arg is specified, the size of the plasma store can be ignored when calculating free memory size.


Discussion:

Resolving patch:
diff --git a/mars/worker/__main__.py b/mars/worker/__main__.py
--- a/mars/worker/__main__.py
+++ b/mars/worker/__main__.py
@@ -47,6 +47,9 @@ def config_args(self, parser):
         parser.add_argument('--cache-mem', help='cache memory size limit')
         parser.add_argument('--min-mem', help='minimal free memory required to start worker')
         parser.add_argument('--spill-dir', help='spill directory')
+        parser.add_argument('--plasma-dir', help='path of plasma directory. When specified, the size '
+                                                 'of plasma store will not be taken into account when '
+                                                 'managing host memory')
 
         compress_types = ', '.join(v.value for v in CompressType.__members__.values())
         parser.add_argument('--disk-compression',
@@ -73,6 +76,8 @@ def validate_arguments(self):
     def create_pool(self, *args, **kwargs):
         # here we create necessary actors on worker
         # and distribute them over processes
+
+        plasma_dir = self.args.plasma_dir or os.environ.get('MARS_PLASMA_DIRS')
         self._service = WorkerService(
             advertise_addr=self.args.advertise,
             n_cpu_process=self.args.cpu_procs,
@@ -84,6 +89,8 @@ def create_pool(self, *args, **kwargs):
             min_mem_size=self.args.min_mem,
             disk_compression=self.args.disk_compression.lower(),
             transfer_compression=self.args.transfer_compression.lower(),
+            plasma_dir=plasma_dir,
+            use_ext_plasma_dir=bool(plasma_dir),
         )
         # start plasma
         self._service.start_plasma()
diff --git a/mars/worker/service.py b/mars/worker/service.py
--- a/mars/worker/service.py
+++ b/mars/worker/service.py
@@ -91,6 +91,9 @@ def __init__(self, **kwargs):
         self._ignore_avail_mem = kwargs.pop('ignore_avail_mem', None) or False
         self._min_mem_size = kwargs.pop('min_mem_size', None) or 128 * 1024 ** 2
 
+        self._plasma_dir = kwargs.pop('plasma_dir', None)
+        self._use_ext_plasma_dir = kwargs.pop('use_ext_plasma_dir', None) or False
+
         self._soft_quota_limit = self._soft_mem_limit
 
         self._calc_memory_limits()
@@ -133,7 +136,8 @@ def _calc_size_limit(limit_str, total_size):
         if self._ignore_avail_mem:
             self._soft_quota_limit = self._soft_mem_limit
         else:
-            self._soft_quota_limit = self._soft_mem_limit - self._cache_mem_limit - actual_used
+            used_cache_size = 0 if self._use_ext_plasma_dir else self._cache_mem_limit
+            self._soft_quota_limit = self._soft_mem_limit - used_cache_size - actual_used
             if self._soft_quota_limit < self._min_mem_size:
                 raise MemoryError('Memory not enough. soft_limit=%s, cache_limit=%s, used=%s' %
                                   tuple(readable_size(k) for k in (
@@ -142,7 +146,8 @@ def _calc_size_limit(limit_str, total_size):
         logger.info('Setting soft limit to %s.', readable_size(self._soft_quota_limit))
 
     def start_plasma(self):
-        self._plasma_store = plasma.start_plasma_store(self._cache_mem_limit)
+        self._plasma_store = plasma.start_plasma_store(
+            self._cache_mem_limit, plasma_directory=self._plasma_dir)
         options.worker.plasma_socket, _ = self._plasma_store.__enter__()
 
     def start(self, endpoint, pool, distributed=True, discoverer=None, process_start_index=0):


Test patch: