Browse Source
!12795 fix GeneratorDataset multithread print error and hung
From: @heleiwang Reviewed-by: Signed-off-by:tags/v1.2.0-rc1
mindspore-ci-bot
Gitee
4 years ago
1 changed files with 4 additions and 3 deletions
Unified View
Diff Options
+ 4
- 3
mindspore/dataset/engine/datasets.py
View File
| @@ -3174,7 +3174,7 @@ class SamplerFn: | |||||
| self.workers = [] | self.workers = [] | ||||
| self.num_worker = num_worker | self.num_worker = num_worker | ||||
| self.multi_process = multi_process | self.multi_process = multi_process | ||||
| self.joined = False | |||||
| self.need_join = False | |||||
| self.ppid = os.getpid() | self.ppid = os.getpid() | ||||
| self.pid = [] | self.pid = [] | ||||
| # Event for end of epoch | # Event for end of epoch | ||||
| @@ -3192,6 +3192,7 @@ class SamplerFn: | |||||
| # In this phase, the main process is not locked. | # In this phase, the main process is not locked. | ||||
| worker.start() | worker.start() | ||||
| self.pid.append(worker.pid) | self.pid.append(worker.pid) | ||||
| self.need_join = True | |||||
| else: | else: | ||||
| worker = _GeneratorWorkerMt(dataset, self.eof) | worker = _GeneratorWorkerMt(dataset, self.eof) | ||||
| worker.daemon = True | worker.daemon = True | ||||
| @@ -3237,9 +3238,9 @@ class SamplerFn: | |||||
| def _stop_subprocess(self): | def _stop_subprocess(self): | ||||
| # Only the main process can call join | # Only the main process can call join | ||||
| if self.joined is False and self.ppid == os.getpid(): | |||||
| if self.need_join is True and self.ppid == os.getpid(): | |||||
| self.eof.set() | self.eof.set() | ||||
| self.joined = True | |||||
| self.need_join = False | |||||
| for w in self.workers: | for w in self.workers: | ||||
| w.join() | w.join() | ||||