mindspore 2.3.0__cp39-cp39-win_amd64.whl → 2.4.0__cp39-cp39-win_amd64.whl

mindspore/dataset/engine/datasets_user_defined.py

@@ -1,4 +1,4 @@
-# Copyright 2019-2023 Huawei Technologies Co., Ltd
+# Copyright 2019-2024 Huawei Technologies Co., Ltd
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -19,12 +19,13 @@ After declaring the dataset object, you can further apply dataset operations
 (e.g. filter, skip, concat, map, batch) on it.
 """
 import builtins
-import copy
 import errno
+import itertools
 import math
 import os
 import signal
 import time
+from types import GeneratorType
 import multiprocessing
 from multiprocessing.util import Finalize
 import queue
@@ -46,7 +47,7 @@ from . import samplers
 from .queue import _SharedQueue
 from .validators import check_generatordataset, check_numpyslicesdataset, check_paddeddataset
 from ..core.config import get_enable_shared_mem, get_prefetch_size, get_multiprocessing_timeout_interval, \
-    get_enable_watchdog, get_debug_mode
+    get_enable_watchdog, get_debug_mode, get_seed, set_seed
 from ..core.datatypes import mstypelist_to_detypelist
 from ..core.py_util_helpers import ExceptionHandler
 from ..transforms import transforms
@@ -89,7 +90,7 @@ def _generator_fn(generator, num_samples):
             yield _convert_row(val)
 
 
-def _cpp_sampler_fn(sample_ids, dataset):
+def _cpp_sampler_fn(dataset, sample_ids):
     """
     Generator function wrapper for mappable dataset with cpp sampler.
     """
@@ -104,7 +105,7 @@ def _cpp_sampler_fn(sample_ids, dataset):
         yield _convert_row(val)
 
 
-def _cpp_sampler_fn_mp(sample_ids, sample_fn):
+def _cpp_sampler_fn_mp(sample_fn, sample_ids):
     """
     Multiprocessing generator function wrapper for mappable dataset with cpp sampler.
     """
@@ -116,6 +117,14 @@ def _cpp_sampler_fn_mp(sample_ids, sample_fn):
     return sample_fn.process(sample_ids)
 
 
+def _generator_fn_wrapper(function, *args):
+    """
+    Generate a new function that wraps the specified generator function with partial
+    application of the given arguments and keywords.
+    """
+    return partial(function, *args)
+
+
 def _fill_worker_indices(workers, indices, idx_cursor, worker_to_quit):
     """
     Worker index queue filler, fill worker index queue in round robin order or QUIT flag.
@@ -178,25 +187,42 @@ def _convert_row(row):
     return tuple(value)
 
 
-class SamplerFn:
+class SamplerFn(cde.PythonMultiprocessingRuntime):
     """
     Multiprocessing or multithread generator function wrapper master process.
     """
 
     def __init__(self, dataset, num_worker, multi_process, max_rowsize):
+        super(SamplerFn, self).__init__()
         self.workers = []
         self.dataset = dataset
         self.num_worker = num_worker
         self.multi_process = multi_process
         self.max_rowsize = max_rowsize
         self.need_join = False
+
+    def is_mp_enabled(self):
+        return self.workers is not None and self.workers
+
+    def launch(self, op_id=-1):
+        """launch the multiprocessing pool"""
+        self.op_id = op_id
+        logger.info("Launching new Python Multiprocessing pool for GeneratorOp:" + str(self.op_id))
+        if self.is_mp_enabled():
+            message = "Launching a new Python multiprocessing pool for GeneratorOp while a pool already exists!" + \
+                " The existing pool will be terminated first."
+            logger.warning(message)
+            self._stop_subprocess()
+            self.reset()
+            self.workers = []
+
         self.ppid = os.getpid()
         self.pids = []
         self.check_interval = get_multiprocessing_timeout_interval()  # the interval of check queue's size
         self._final_join = True
 
         # Event for end of epoch
-        if multi_process is True:
+        if self.multi_process is True:
             try:
                 self.eof = multiprocessing.Event()
             except Exception:
@@ -206,22 +232,22 @@ class SamplerFn:
             self.eof = threading.Event()
         # Create workers
 
-        # get default queue size and adjust queuesize per worker if there are large # workers
+        # get default queue size and adjust queue size per worker if there are large # workers
         queue_size = get_prefetch_size()
-        queue_size = min(queue_size, queue_size * 4 // num_worker)
+        queue_size = min(queue_size, queue_size * 4 // self.num_worker)
         queue_size = max(2, queue_size)
 
-        if multi_process and get_enable_shared_mem():
+        if self.multi_process and get_enable_shared_mem():
             # generator dataset use idx_queue and res_queue to transfer data between main and subprocess
             # idx_queue is used multiprocess.Queue which is not shared memory, so it's size is 0.
-            # res_queue is used shared memory, so it' size is max_rowsize which is defined by user.
-            _check_shm_usage(num_worker, queue_size, 0, max_rowsize)
+            # res_queue is used shared memory, so its size is max_rowsize which is defined by user.
+            _check_shm_usage(self.num_worker, queue_size, 0, self.max_rowsize)
         self.count = multiprocessing.Value('i', 0)
-        for worker_id in range(num_worker):
-            if multi_process is True:
+        for worker_id in range(self.num_worker):
+            if self.multi_process is True:
                 try:
-                    worker = _GeneratorWorkerMp(dataset, self.eof, max_rowsize, queue_size, self.ppid, self.count,
-                                                worker_id)
+                    worker = _GeneratorWorkerMp(self.dataset, self.eof, self.max_rowsize, queue_size, self.ppid,
+                                                self.count, worker_id)
                     worker.daemon = True
                     # When multi processes fork a subprocess, the lock of the main process is copied to the subprocess,
                     # which may cause deadlock. Therefore, the subprocess startup is performed in the initialization
@@ -240,10 +266,12 @@ class SamplerFn:
                 self.pids.append(worker.pid)
                 self.need_join = True
             else:
-                worker = _GeneratorWorkerMt(dataset, self.eof, worker_id)
+                worker = _GeneratorWorkerMt(self.dataset, self.eof, worker_id)
                 worker.daemon = True
+                self.need_join = True
             self.workers.append(worker)
-        self._launch_cleanup_worker(multi_process=multi_process)
+        if self.multi_process and platform.system().lower() != 'windows':
+            self._launch_cleanup_worker()
 
     def _interval_log(self, i, start_time, wait_count):
         cost_time = int(time.time()) - start_time
@@ -252,11 +280,10 @@ class SamplerFn:
             self._log_stuck_warning(self.workers[i % self.num_worker], cost_time)
         return wait_count
 
-    def process(self, indices):
-        """
-        The main process, start the child process or child thread, and fill the index queue.
-        Get the result and return.
-        """
+    def _check_and_start_process(self):
+        """Check the idx_queue and start the process"""
+        if self.workers is None:
+            raise RuntimeError("The GeneratorDataset subprocess worker may be killed or exit abnormally.")
         for w in self.workers:
             # Check whether the queue of the subprocess is empty.
             if not w.queue_empty():
@@ -270,7 +297,20 @@ class SamplerFn:
                         continue
             # Start all workers
             if not w.is_alive():
-                w.start()
+                try:
+                    w.start()
+                except RuntimeError as e:
+                    # the worker may be being started.
+                    if w._started.is_set():  # pylint: disable=W0212
+                        continue
+                    raise e
+
+    def process(self, indices):
+        """
+        The main process, start the child process or child thread, and fill the index queue.
+        Get the result and return.
+        """
+        self._check_and_start_process()
 
         # Fill initial index queues
         idx_cursor = 0
@@ -300,14 +340,6 @@ class SamplerFn:
                     time.sleep(0.1)
                     wait_count = self._interval_log(i, start_time, wait_count)
                 result = self.workers[i % self.num_worker].get()
-                # Because there is no need to copy when creating Tensors in the C++layer, it reduces the time
-                # from np.ndarray to C++Tensor creation. However, when using shared memory in multiple processes,
-                # the address of the shared memory will always be passed to subsequent nodes in the dataset pipeline,
-                # and the shared memory will also be written by the current node, causing dirty data to be accessed
-                # by subsequent nodes in the pipeline. So make a memory copy here to solve the problem of
-                # shared memory being contaminated.
-                if self.multi_process is True and get_enable_shared_mem():
-                    result = copy.deepcopy(result)
                 if isinstance(result, ExceptionHandler):
                     result.reraise()
             except queue.Empty:
@@ -360,44 +392,74 @@ class SamplerFn:
                            "the `mindspore.dataset.config.set_multiprocessing_timeout_interval` interface."
         logger.warning(warning_message)
 
-    def _launch_cleanup_worker(self, multi_process):
+    def _launch_cleanup_worker(self):
         """
         We need a extra thread and process if main process or subprocess was killed.
-
-        Args:
-            multi_process: Whether use multiprocess.
         """
-        if multi_process is True and platform.system().lower() != 'windows':
-            _clean_worker_func = _PythonMultiprocessing._clean_process  # pylint: disable=W0212
-            self.cleaning_process = multiprocessing.Process(target=_clean_worker_func,
-                                                            name="GeneratorCleanProcess",
-                                                            args=(self.ppid, self.workers, self.eof))
-            self.cleaning_process.daemon = True
-            self.cleaning_process.start()
-
-            if get_enable_watchdog():
-                self.eot = threading.Event()
-                self.watch_dog = threading.Thread(target=_PythonMultiprocessing._watch_dog,  # pylint: disable=W0212
-                                                  name="GeneratorWatchDog",
-                                                  args=(self.eot, self.workers + [self.cleaning_process]))
-                self.watch_dog.daemon = True
-                self.watch_dog.start()
-
-                if self._final_join is True:
-                    self._jointhread = Finalize(
-                        self.watch_dog, self._finalize_join,
-                        args=(weakref.ref(self.watch_dog), self.eot),
-                        exitpriority=-5
-                    )
+        _clean_worker_func = _PythonMultiprocessing._clean_process  # pylint: disable=W0212
+        self.cleaning_process = multiprocessing.Process(target=_clean_worker_func,
+                                                        name="GeneratorCleanProcess",
+                                                        args=(self.ppid, self.workers, self.eof))
+        self.cleaning_process.daemon = True
+        self.cleaning_process.start()
+
+        if get_enable_watchdog():
+            self.eot = threading.Event()
+            self.watch_dog = threading.Thread(target=_PythonMultiprocessing._watch_dog,  # pylint: disable=W0212
+                                              name="GeneratorWatchDog",
+                                              args=(self.eot, self.workers + [self.cleaning_process]))
+            self.watch_dog.daemon = True
+            self.watch_dog.start()
+
+            if self._final_join is True:
+                self._jointhread = Finalize(
+                    self.watch_dog, self._finalize_join,
+                    args=(weakref.ref(self.watch_dog), self.eot),
+                    exitpriority=-5
+                )
+
+    def _release_fd(self):
+        """Release the file descriptor by subprocess"""
+        # release the file descriptor handle
+        check_interval = get_multiprocessing_timeout_interval()
+        for w in self.workers:
+            try:
+                subprocess_file_descriptor = w.sentinel
+                st = time.time()
+                while _PythonMultiprocessing.is_process_alive(w.pid):
+                    time.sleep(0.01)  # sleep 10ms, waiting for the subprocess exit
+                    if time.time() - st > check_interval:
+                        logger.warning("Waiting for the subprocess worker [{}] to exit.".format(w.pid))
+                        st += check_interval
+            except ValueError as e:
+                if "process object is closed" in str(e):
+                    continue
+                raise e
+            try:
+                if w.is_alive():
+                    os.close(subprocess_file_descriptor)
+            except OSError as e:
+                # Maybe the file descriptor had been released, so ignore the 'Bad file descriptor'
+                if "Bad file descriptor" not in str(e):
+                    raise e
+            except AttributeError:  # maybe occur "'NoneType' object has no attribute 'maxsize'"
+                pass
 
     def _stop_subprocess(self):
-        """Only the main process can call join."""
+        """Only the main process can call join. All the sub-process / sub-thread will be stopped."""
         if self.need_join is True and self.ppid == os.getpid():
+            # the sub-process / sub-thread will stop by self.eof.set()
             if hasattr(self, 'eof') and self.eof is not None:
-                self.eof.set()
+                try:
+                    self.eof.set()
+                except AttributeError:  # maybe occur "'NoneType' object has no attribute 'maxsize'"
+                    pass
+
             # close the watch dog first
             self._abort_watchdog()
             self.need_join = False
+
+            # waiting for the sub-process stop
             for w in self.workers:
                 if self.multi_process is True and hasattr(w, '_closed') and w._closed is False:  # pylint: disable=W0212
                     try:
@@ -415,28 +477,8 @@ class SamplerFn:
                         # Block all errors when join
                         continue
 
-            # release the file descriptor handle
-            check_interval = get_multiprocessing_timeout_interval()
-            for w in self.workers:
-                try:
-                    subprocess_file_descriptor = w.sentinel
-                    st = time.time()
-                    while _PythonMultiprocessing.is_process_alive(w.pid):
-                        time.sleep(0.01)  # sleep 10ms, waiting for the subprocess exit
-                        if time.time() - st > check_interval:
-                            logger.warning("Waiting for the subprocess worker [{}] to exit.".format(w.pid))
-                            st += check_interval
-                except ValueError as e:
-                    if "process object is closed" in str(e):
-                        continue
-                    raise e
-                try:
-                    if w.is_alive():
-                        os.close(subprocess_file_descriptor)
-                except OSError as e:
-                    # Maybe the file descriptor had been released, so ignore the 'Bad file descriptor'
-                    if "Bad file descriptor" not in str(e):
-                        raise e
+            if self.multi_process is True:
+                self._release_fd()
 
             self.workers.clear()
             self.workers = None
@@ -498,13 +540,21 @@ def _main_process_already_exit(eof, is_multiprocessing, idx_queue, result_queue,
     return False
 
 
-def _generator_worker_loop(dataset, idx_queue, result_queue, eof, is_multiprocessing, ppid=-1):
+def _generator_worker_loop(dataset, idx_queue, result_queue, eof, is_multiprocessing, worker_id, ppid=-1):
     """
     Multithread or multiprocess generator worker process loop.
     """
+    # Initialize C++ side signal handlers
+    cde.register_worker_handlers()
+
     if is_multiprocessing:
         result_queue.cancel_join_thread()  # Ensure that the process does not hung when exiting
         signal.signal(signal.SIGTERM, partial(_subprocess_handle, eof))
+
+        # init the random seed and np.random seed for the subprocess
+        if get_seed() != 5489:
+            set_seed(get_seed() + worker_id)
+
     while not eof.is_set():
         _ignore_sigint(is_multiprocessing=is_multiprocessing)
 
@@ -562,7 +612,8 @@ class _GeneratorWorkerMt(threading.Thread):
     def __init__(self, dataset, eof, worker_id):
         self.idx_queue = queue.Queue(16)
         self.res_queue = queue.Queue(16)
-        super().__init__(target=_generator_worker_loop, args=(dataset, self.idx_queue, self.res_queue, eof, False),
+        super().__init__(target=_generator_worker_loop,
+                         args=(dataset, self.idx_queue, self.res_queue, eof, False, worker_id),
                          name="GeneratorWorkerThread" + str(worker_id))
 
     def put(self, item):
@@ -598,8 +649,9 @@ class _GeneratorWorkerMp(multiprocessing.Process):
             self.res_queue = _SharedQueue(queue_size, count, max_rowsize=max_rowsize)
         else:
             self.res_queue = multiprocessing.Queue(queue_size)
-        self.idx_queue.cancel_join_thread()  # Ensure that the process does not hung when exiting
-        super().__init__(target=_generator_worker_loop, args=(dataset, self.idx_queue, self.res_queue, eof, True, ppid),
+        self.idx_queue.cancel_join_thread()  # Ensure that the process does not hang when exiting
+        super().__init__(target=_generator_worker_loop,
+                         args=(dataset, self.idx_queue, self.res_queue, eof, True, worker_id, ppid),
                          name="GeneratorWorkerProcess" + str(worker_id))
 
     def put(self, item):
@@ -634,6 +686,20 @@ class _GeneratorWorkerMp(multiprocessing.Process):
             del self.res_queue
 
 
+class _GeneratorWrapper:
+    """Wrapper the generator so that it can be iterated multiple times in GeneratorDataset."""
+    def __init__(self, generator):
+        self.generator = generator
+        self.generator_new, self.generator = itertools.tee(self.generator)
+
+    def __iter__(self):
+        self.generator_new, self.generator = itertools.tee(self.generator)
+        return self
+
+    def __next__(self):
+        return next(self.generator_new)
+
+
 class GeneratorDataset(MappableDataset, UnionBaseDataset):
     """
     A source dataset that generates data from Python by invoking Python data source each epoch.
@@ -671,11 +737,11 @@ class GeneratorDataset(MappableDataset, UnionBaseDataset):
             Random accessible input is required.
         python_multiprocessing (bool, optional): Parallelize Python operations with multiple worker process. This
             option could be beneficial if the Python operation is computational heavy. Default: ``True``.
-        max_rowsize(int, optional): Maximum size of row in MB that is used for shared memory
+        max_rowsize(int, optional): Maximum size of data (in MB) that is used for shared memory
             allocation to copy data between processes, the total occupied shared memory will increase as
             ``num_parallel_workers`` and :func:`mindspore.dataset.config.set_prefetch_size` increase. If set to -1,
             shared memory will be dynamically allocated with the actual size of data. This is only used if
-            ``python_multiprocessing`` is set to True. Default: 16.
+            ``python_multiprocessing`` is set to True. Default: ``None`` , allocate shared memory dynamically.
 
     Raises:
         RuntimeError: If source raises an exception during execution.
@@ -693,16 +759,16 @@ class GeneratorDataset(MappableDataset, UnionBaseDataset):
 
     Note:
         - If you configure `python_multiprocessing=True` (Default: ``True`` ) and `num_parallel_workers>1`
-          (default: ``1`` ) indicates that the multi-process mode is started for data load acceleration.
+          (default: ``1`` ) indicates that the multiprocessing mode is started for data load acceleration.
           At this time, as the datasetiterates, the memory consumption of the subprocess will gradually increase,
           mainly because the subprocess of the user-defined dataset obtains the member variables from the main
           process in the Copy On Write way.
           Example: If you define a dataset with `__ init__` function which contains a large number of member variable
           data (for example, a very large file name list is loaded during the dataset construction) and uses the
-          multi-process mode, which may cause the problem of OOM (the estimated total memory usage is:
+          multiprocessing mode, which may cause the problem of OOM (the estimated total memory usage is:
           `(num_parallel_workers+1) * size of the parent process` ). The simplest solution is to replace Python objects
           (such as list/dict/int/float/string) with non referenced data types
-          (such as Pandas, Numpy or PyArrow objects) for member variables, or load less meta data in member variables,
+          (such as Pandas, Numpy or PyArrow objects) for member variables, or load less metadata in member variables,
           or configure `python_multiprocessing=False` to use multi-threading mode.
 
           There are several classes/functions that can help you reduce the size of member variables, and you can choose
@@ -782,7 +848,7 @@ class GeneratorDataset(MappableDataset, UnionBaseDataset):
     @check_generatordataset
     def __init__(self, source, column_names=None, column_types=None, schema=None, num_samples=None,
                  num_parallel_workers=1, shuffle=None, sampler=None, num_shards=None, shard_id=None,
-                 python_multiprocessing=True, max_rowsize=6):
+                 python_multiprocessing=True, max_rowsize=None):
         super().__init__(num_parallel_workers=num_parallel_workers, sampler=sampler, num_samples=num_samples,
                          shuffle=shuffle, num_shards=num_shards, shard_id=shard_id)
         if isinstance(source, builtins.zip):
@@ -790,6 +856,11 @@ class GeneratorDataset(MappableDataset, UnionBaseDataset):
             self.source = [item for item in source]
         else:
             self.source = source
+
+        # wrapper the generator so that it can be iterated multiple times
+        if isinstance(self.source, GeneratorType):
+            self.source = _GeneratorWrapper(self.source)
+
         self.prepared_source = None  # source to be sent to C++
         if hasattr(self, 'operator_mixed') and getattr(self, 'operator_mixed') is True:
             self.num_parallel_workers = 1
@@ -805,7 +876,6 @@ class GeneratorDataset(MappableDataset, UnionBaseDataset):
         if self.python_multiprocessing and get_debug_mode():
             logger.warning("Python multiprocessing is not supported in debug mode."
                            " Ignoring Python multiprocessing for GeneratorDataset.")
-            self.python_multiprocessing = False
 
         self.column_names = to_list(column_names)
 
@@ -829,7 +899,7 @@ class GeneratorDataset(MappableDataset, UnionBaseDataset):
             if isinstance(self.sampler, samplers.Sampler) or hasattr(self.sampler, "__iter__"):
                 self.source_len = len(list(sampler))
 
-        self.max_rowsize = max_rowsize
+        self.max_rowsize = max_rowsize if max_rowsize is not None else -1
         self.sample_fn = None
 
     def __deepcopy__(self, memodict):
@@ -863,14 +933,14 @@ class GeneratorDataset(MappableDataset, UnionBaseDataset):
             if self.source_len == -1:
                 raise RuntimeError("Attempt to construct a random access dataset, '__len__' method is required!")
 
-            if self.num_parallel_workers > 1:
+            if self.num_parallel_workers > 1 and not get_debug_mode():
                 self.__validate_memory_usage()
 
                 sample_fn = SamplerFn(self.source, self.num_parallel_workers, self.python_multiprocessing,
                                       self.max_rowsize)
-                self.prepared_source = (lambda sample_ids: _cpp_sampler_fn_mp(sample_ids, sample_fn))
+                self.prepared_source = _generator_fn_wrapper(_cpp_sampler_fn_mp, sample_fn)
             else:
-                self.prepared_source = (lambda sample_ids: _cpp_sampler_fn(sample_ids, self.source))
+                self.prepared_source = _generator_fn_wrapper(_cpp_sampler_fn, self.source)
             self.sample_fn = sample_fn
         else:
             self.sampler = None
@@ -878,30 +948,30 @@ class GeneratorDataset(MappableDataset, UnionBaseDataset):
             self.source_len = min(self.source_len, self.num_samples) if self.num_samples != 0 else self.source_len
             if not hasattr(self.source, "__iter__"):
                 # Use generator function if input callable
-                self.prepared_source = (lambda: _generator_fn(self.source, self.num_samples))
+                self.prepared_source = _generator_fn_wrapper(_generator_fn, self.source, self.num_samples)
             else:
                 # Use iterator function if input is iterable
                 # Random accessible input is also iterable
-                self.prepared_source = (lambda: _iter_fn(self.source, self.num_samples))
+                self.prepared_source = _generator_fn_wrapper(_iter_fn, self.source, self.num_samples)
 
     def parse(self, children=None):
         self.prepare_multiprocessing()
         if self.schema is None:
             return cde.GeneratorNode(self.prepared_source, self.column_names, self.column_types, self.source_len,
-                                     self.sampler, self.num_parallel_workers)
+                                     self.sampler, self.num_parallel_workers, self.sample_fn)
         schema = self.schema
         if isinstance(schema, Schema):
             schema = self.schema.cpp_schema
         return cde.GeneratorNode(self.prepared_source, schema, self.source_len, self.sampler,
-                                 self.num_parallel_workers)
+                                 self.num_parallel_workers, self.sample_fn)
 
     def __validate_memory_usage(self):
         """
-        Check memory usage when mulit-processing mode, when 85% prompt warning and 100% raise error.
+        Check memory usage when multiprocessing mode, when 85% prompt warning and 100% raise error.
         """
         if self.python_multiprocessing:
-            # if use num_parallel_workers is to large when python_multiprocessing=True which would cause
-            # OOM error get the num_shards
+            # setting num_parallel_workers too large when using python multiprocessing may cause
+            # out of memory for getting num_shards
             valid_num_shards = 1
             if isinstance(self.sampler, samplers.DistributedSampler):
                 valid_num_shards = self.sampler.num_shards