Create the Distributed Strategy session.
(distribution_strategy)
| 5785 | """Configure session config and create a session with it.""" |
| 5786 | |
| 5787 | def _create_session(distribution_strategy): |
| 5788 | """Create the Distributed Strategy session.""" |
| 5789 | session_config = get_default_session_config() |
| 5790 | |
| 5791 | # If a session already exists, merge in its config; in the case there is a |
| 5792 | # conflict, take values of the existing config. |
| 5793 | global _SESSION |
| 5794 | if getattr(_SESSION, 'session', None) and _SESSION.session._config: |
| 5795 | session_config.MergeFrom(_SESSION.session._config) |
| 5796 | |
| 5797 | if is_tpu_strategy(distribution_strategy): |
| 5798 | # TODO(priyag, yuefengz): Remove this workaround when Distribute |
| 5799 | # Coordinator is integrated with keras and we can create a session from |
| 5800 | # there. |
| 5801 | distribution_strategy.configure(session_config) |
| 5802 | master = distribution_strategy.extended._tpu_cluster_resolver.master() # pylint: disable=protected-access |
| 5803 | session = session_module.Session(config=session_config, target=master) |
| 5804 | else: |
| 5805 | worker_context = dc_context.get_current_worker_context() |
| 5806 | if worker_context: |
| 5807 | dc_session_config = worker_context.session_config |
| 5808 | # Merge the default session config to the one from distribute |
| 5809 | # coordinator, which is fine for now since they don't have |
| 5810 | # conflicting configurations. |
| 5811 | dc_session_config.MergeFrom(session_config) |
| 5812 | session = session_module.Session( |
| 5813 | config=dc_session_config, target=worker_context.master_target) |
| 5814 | else: |
| 5815 | distribution_strategy.configure(session_config) |
| 5816 | session = session_module.Session(config=session_config) |
| 5817 | |
| 5818 | set_session(session) |
| 5819 | |
| 5820 | if distribution_strategy._in_multi_worker_mode(): |
| 5821 | dc.run_distribute_coordinator( |
no test coverage detected