# Copyright 2017 The TensorFlow Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== """Environment configuration object for Estimators.""" from __future__ import absolute_import from __future__ import division from __future__ import print_function import copy import json import os import six import tensorflow as tf from tensorflow.core.protobuf import rewriter_config_pb2 from tensorflow.python.distribute import estimator_training as distribute_coordinator_training from tensorflow.python.util import function_utils from tensorflow_estimator.python.estimator.estimator_export import estimator_export _USE_DEFAULT = object() _VALID_DEVICE_FN_ARGS = set(['op']) # A list of the property names in RunConfig that the user is allowed to change. _DEFAULT_REPLACEABLE_LIST = [ 'model_dir', 'tf_random_seed', 'save_summary_steps', 'save_checkpoints_steps', 'save_checkpoints_secs', 'session_config', 'keep_checkpoint_max', 'keep_checkpoint_every_n_hours', 'log_step_count_steps', 'train_distribute', 'device_fn', 'protocol', 'eval_distribute', 'experimental_distribute', 'experimental_max_worker_delay_secs', 'session_creation_timeout_secs', 'checkpoint_save_graph_def' ] _SAVE_CKPT_ERR = ( '`save_checkpoints_steps` and `save_checkpoints_secs` cannot be both set.') _TF_CONFIG_ENV = 'TF_CONFIG' _TASK_ENV_KEY = 'task' _TASK_TYPE_KEY = 'type' _TASK_ID_KEY = 'index' _CLUSTER_KEY = 'cluster' _SERVICE_KEY = 'service' _SESSION_MASTER_KEY = 'session_master' _EVAL_SESSION_MASTER_KEY = 'eval_session_master' _MODEL_DIR_KEY = 'model_dir' _LOCAL_MASTER = '' _GRPC_SCHEME = 'grpc://' def _get_session_master(cluster_spec, task_type, task_id, tf_config): """Returns the appropriate address for TensorFlow master. The order of precedence to determine the TF session master is as follows: 1. If `tf_session_master` is set in TF_CONFIG environment variable, takes it. 2. If the cluster has only one node, returns empty string ''. 3. Returns the grpc address according to the task type and id in the cluster. This is between-graph replication. Note: task_type and task_id must be validated. Typically, validated using `_validate_task_type_and_task_id`. Args: cluster_spec: A `ClusterSpec` instance. task_type: String. Task type for current node. task_id: Int. Task id for current node. tf_config: Dict. Python dict for the TF_CONFIG environment variable. Raises: RuntimeError: If `cluster_spec` is not set. """ if _SESSION_MASTER_KEY in tf_config: return tf_config[_SESSION_MASTER_KEY] if not cluster_spec: raise RuntimeError('Internal error: `_get_session_master` ' 'does not expect empty cluster_spec.') jobs = cluster_spec.jobs # If there is only one node in the cluster, do things locally by setting # master to ''. If a service or user sets TF_CONFIG with a single node, it's # more performant to use a direct master rather than an RPC service. if len(jobs) == 1 and len(cluster_spec.job_tasks(jobs[0])) == 1: return _LOCAL_MASTER # Lookup the master in cluster_spec using task_type and task_id, # if possible. addresses = cluster_spec.job_tasks(task_type) return _GRPC_SCHEME + addresses[task_id] def _get_eval_session_master(task_type, tf_config): """Returns the appropriate address for TensorFlow evaluation master.""" if task_type == TaskType.EVALUATOR: return tf_config.get(_EVAL_SESSION_MASTER_KEY, _LOCAL_MASTER) return _LOCAL_MASTER def _count_ps(cluster_spec): """Counts the number of parameter servers in cluster_spec.""" if not cluster_spec: raise RuntimeError( 'Internal error: `_count_ps` does not expect empty cluster_spec.') return len(cluster_spec.as_dict().get(TaskType.PS, [])) def _count_worker(cluster_spec, chief_task_type): """Counts the number of workers (including chief) in cluster_spec.""" if not cluster_spec: raise RuntimeError( 'Internal error: `_count_worker` does not expect empty cluster_spec.') return (len(cluster_spec.as_dict().get(TaskType.WORKER, [])) + len(cluster_spec.as_dict().get(chief_task_type, []))) def _validate_service(service): """Validates the service key.""" if service is not None and not isinstance(service, dict): raise TypeError( 'If "service" is set in TF_CONFIG, it must be a dict. Given %s' % type(service)) return service def _validate_task_type_and_task_id(cluster_spec, task_env, chief_task_type): """Validates the task type and index in `task_env` according to cluster.""" if chief_task_type not in cluster_spec.jobs: raise ValueError( 'If "cluster" is set in TF_CONFIG, it must have one "%s" node.' % chief_task_type) if len(cluster_spec.job_tasks(chief_task_type)) > 1: raise ValueError( 'The "cluster" in TF_CONFIG must have only one "%s" node.' % chief_task_type) task_type = task_env.get(_TASK_TYPE_KEY, None) task_id = task_env.get(_TASK_ID_KEY, None) if not task_type: raise ValueError('If "cluster" is set in TF_CONFIG, task type must be set.') if task_id is None: raise ValueError( 'If "cluster" is set in TF_CONFIG, task index must be set.') task_id = int(task_id) # Check the task id bounds. Upper bound is not necessary as # - for evaluator, there is no upper bound. # - for non-evaluator, task id is upper bounded by the number of jobs in # cluster spec, which will be checked later (when retrieving the `master`) if task_id < 0: raise ValueError('Task index must be non-negative number.') # Evaluator is not part of the training cluster. if task_type == TaskType.EVALUATOR: return task_type, task_id if task_type not in cluster_spec.jobs: raise ValueError( '%s is not a valid task_type in the cluster_spec:\n' '%s\n\n' 'Note that these values may be coming from the TF_CONFIG environment ' 'variable.' % (task_type, cluster_spec)) addresses = cluster_spec.job_tasks(task_type) if not 0 <= task_id < len(addresses): raise ValueError( '%d is not a valid task_id for task_type %s in the cluster_spec:\n' '%s\n\n' 'Note that these values may be coming from the TF_CONFIG environment ' 'variable.' % (task_id, task_type, cluster_spec)) return task_type, task_id def _get_global_id_in_cluster(cluster_spec, task_type, task_id, chief_task_type): """Returns the global id in cluster.""" # Note: This is implementation details, which user should not rely on. # The first id is 0, which is always for the `chief` node. All other nodes, # except `ps`, are ordered alphabetical based on task type (alphabetically) # and task id (ascendingly). `ps` are ordered last. # Sort task names in cluster task_type_ordered_list = [chief_task_type] task_type_ordered_list.extend([ t for t in sorted(cluster_spec.jobs) if t != chief_task_type and t != TaskType.PS ]) if TaskType.PS in cluster_spec.jobs: task_type_ordered_list.append(TaskType.PS) next_global_id = 0 for t in task_type_ordered_list: if t == task_type: return next_global_id + task_id next_global_id += len(cluster_spec.job_tasks(t)) # This should never happen. raise RuntimeError('Internal Error: `task_type` ({}) is not in ' 'cluster_spec ({}).'.format(task_type, cluster_spec)) def _validate_save_ckpt_with_replaced_keys(new_copy, replaced_keys): """Validates the save ckpt properties.""" # Ensure one (and only one) of save_steps and save_secs is not None. # Also, if user sets one save ckpt property, say steps, the other one (secs) # should be set as None to improve usability. save_steps = new_copy.save_checkpoints_steps save_secs = new_copy.save_checkpoints_secs if ('save_checkpoints_steps' in replaced_keys and 'save_checkpoints_secs' in replaced_keys): # If user sets both properties explicitly, we need to error out if both # are set or neither of them are set. if save_steps is not None and save_secs is not None: raise ValueError(_SAVE_CKPT_ERR) elif 'save_checkpoints_steps' in replaced_keys and save_steps is not None: new_copy._save_checkpoints_secs = None # pylint: disable=protected-access elif 'save_checkpoints_secs' in replaced_keys and save_secs is not None: new_copy._save_checkpoints_steps = None # pylint: disable=protected-access def _validate_properties(run_config): """Validates the properties.""" def _validate(property_name, cond, message): property_value = getattr(run_config, property_name) if property_value is not None and not cond(property_value): raise ValueError(message) def _validate_delay(delay): """Check that delay is an integer value. Since this has to work for both Python2 and Python3 and PEP237 defines long to be basically int, we cannot just use a lambda function. """ try: return isinstance(delay, (int, long)) except NameError: # PEP237 redefines long to int for Python3 return isinstance(delay, int) _validate( 'model_dir', lambda dir: dir, message='model_dir should be non-empty') _validate( 'save_summary_steps', lambda steps: steps >= 0, message='save_summary_steps should be >= 0') _validate( 'save_checkpoints_steps', lambda steps: steps >= 0, message='save_checkpoints_steps should be >= 0') _validate( 'save_checkpoints_secs', lambda secs: secs >= 0, message='save_checkpoints_secs should be >= 0') _validate( 'session_config', lambda sc: isinstance(sc, tf.compat.v1.ConfigProto), message='session_config must be instance of ConfigProto') _validate( 'keep_checkpoint_max', lambda keep_max: keep_max >= 0, message='keep_checkpoint_max should be >= 0') _validate( 'keep_checkpoint_every_n_hours', lambda keep_hours: keep_hours > 0, message='keep_checkpoint_every_n_hours should be > 0') _validate( 'log_step_count_steps', lambda num_steps: num_steps > 0, message='log_step_count_steps should be > 0') _validate( 'tf_random_seed', lambda seed: isinstance(seed, six.integer_types), message='tf_random_seed must be integer.') _validate( 'experimental_max_worker_delay_secs', _validate_delay, message='experimental_max_worker_delay_secs must be an integer if' ' set.') _validate( 'session_creation_timeout_secs', lambda timeout_secs: timeout_secs > 0, message='session_creation_timeout_secs should be > 0') _validate( 'device_fn', lambda device_fn: six.callable(device_fn) and set( function_utils.fn_args(device_fn)) == _VALID_DEVICE_FN_ARGS, message='device_fn must be callable with exactly' ' one argument "op".') _validate( 'protocol', lambda protocol: protocol in (None, 'grpc', 'grpc+verbs'), message='protocol should be grpc or grpc+verbs') def get_default_session_config(): """Returns tf.ConfigProto instance.""" rewrite_opts = rewriter_config_pb2.RewriterConfig( meta_optimizer_iterations=rewriter_config_pb2.RewriterConfig.ONE) graph_opts = tf.compat.v1.GraphOptions(rewrite_options=rewrite_opts) return tf.compat.v1.ConfigProto( allow_soft_placement=True, graph_options=graph_opts) class TaskType(object): MASTER = 'master' PS = 'ps' WORKER = 'worker' CHIEF = 'chief' EVALUATOR = 'evaluator' @estimator_export('estimator.RunConfig') class RunConfig(object): """This class specifies the configurations for an `Estimator` run.""" def __init__(self, model_dir=None, tf_random_seed=None, save_summary_steps=100, save_checkpoints_steps=_USE_DEFAULT, save_checkpoints_secs=_USE_DEFAULT, session_config=None, keep_checkpoint_max=5, keep_checkpoint_every_n_hours=10000, log_step_count_steps=100, train_distribute=None, device_fn=None, protocol=None, eval_distribute=None, experimental_distribute=None, experimental_max_worker_delay_secs=None, session_creation_timeout_secs=7200, checkpoint_save_graph_def=True): """Constructs a RunConfig. All distributed training related properties `cluster_spec`, `is_chief`, `master` , `num_worker_replicas`, `num_ps_replicas`, `task_id`, and `task_type` are set based on the `TF_CONFIG` environment variable, if the pertinent information is present. The `TF_CONFIG` environment variable is a JSON object with attributes: `cluster` and `task`. `cluster` is a JSON serialized version of `ClusterSpec`'s Python dict from `server_lib.py`, mapping task types (usually one of the `TaskType` enums) to a list of task addresses. `task` has two attributes: `type` and `index`, where `type` can be any of the task types in `cluster`. When `TF_CONFIG` contains said information, the following properties are set on this class: * `cluster_spec` is parsed from `TF_CONFIG['cluster']`. Defaults to {}. If present, must have one and only one node in the `chief` attribute of `cluster_spec`. * `task_type` is set to `TF_CONFIG['task']['type']`. Must set if `cluster_spec` is present; must be `worker` (the default value) if `cluster_spec` is not set. * `task_id` is set to `TF_CONFIG['task']['index']`. Must set if `cluster_spec` is present; must be 0 (the default value) if `cluster_spec` is not set. * `master` is determined by looking up `task_type` and `task_id` in the `cluster_spec`. Defaults to ''. * `num_ps_replicas` is set by counting the number of nodes listed in the `ps` attribute of `cluster_spec`. Defaults to 0. * `num_worker_replicas` is set by counting the number of nodes listed in the `worker` and `chief` attributes of `cluster_spec`. Defaults to 1. * `is_chief` is determined based on `task_type` and `cluster`. There is a special node with `task_type` as `evaluator`, which is not part of the (training) `cluster_spec`. It handles the distributed evaluation job. Example of non-chief node: ``` cluster = {'chief': ['host0:2222'], 'ps': ['host1:2222', 'host2:2222'], 'worker': ['host3:2222', 'host4:2222', 'host5:2222']} os.environ['TF_CONFIG'] = json.dumps( {'cluster': cluster, 'task': {'type': 'worker', 'index': 1}}) config = RunConfig() assert config.master == 'host4:2222' assert config.task_id == 1 assert config.num_ps_replicas == 2 assert config.num_worker_replicas == 4 assert config.cluster_spec == server_lib.ClusterSpec(cluster) assert config.task_type == 'worker' assert not config.is_chief ``` Example of chief node: ``` cluster = {'chief': ['host0:2222'], 'ps': ['host1:2222', 'host2:2222'], 'worker': ['host3:2222', 'host4:2222', 'host5:2222']} os.environ['TF_CONFIG'] = json.dumps( {'cluster': cluster, 'task': {'type': 'chief', 'index': 0}}) config = RunConfig() assert config.master == 'host0:2222' assert config.task_id == 0 assert config.num_ps_replicas == 2 assert config.num_worker_replicas == 4 assert config.cluster_spec == server_lib.ClusterSpec(cluster) assert config.task_type == 'chief' assert config.is_chief ``` Example of evaluator node (evaluator is not part of training cluster): ``` cluster = {'chief': ['host0:2222'], 'ps': ['host1:2222', 'host2:2222'], 'worker': ['host3:2222', 'host4:2222', 'host5:2222']} os.environ['TF_CONFIG'] = json.dumps( {'cluster': cluster, 'task': {'type': 'evaluator', 'index': 0}}) config = RunConfig() assert config.master == '' assert config.evaluator_master == '' assert config.task_id == 0 assert config.num_ps_replicas == 0 assert config.num_worker_replicas == 0 assert config.cluster_spec == {} assert config.task_type == 'evaluator' assert not config.is_chief ``` N.B.: If `save_checkpoints_steps` or `save_checkpoints_secs` is set, `keep_checkpoint_max` might need to be adjusted accordingly, especially in distributed training. For example, setting `save_checkpoints_secs` as 60 without adjusting `keep_checkpoint_max` (defaults to 5) leads to situation that checkpoint would be garbage collected after 5 minutes. In distributed training, the evaluation job starts asynchronously and might fail to load or find the checkpoint due to race condition. Args: model_dir: directory where model parameters, graph, etc are saved. If `PathLike` object, the path will be resolved. If `None`, will use a default value set by the Estimator. tf_random_seed: Random seed for TensorFlow initializers. Setting this value allows consistency between reruns. save_summary_steps: Save summaries every this many steps. save_checkpoints_steps: Save checkpoints every this many steps. Can not be specified with `save_checkpoints_secs`. save_checkpoints_secs: Save checkpoints every this many seconds. Can not be specified with `save_checkpoints_steps`. Defaults to 600 seconds if both `save_checkpoints_steps` and `save_checkpoints_secs` are not set in constructor. If both `save_checkpoints_steps` and `save_checkpoints_secs` are `None`, then checkpoints are disabled. session_config: a ConfigProto used to set session parameters, or `None`. keep_checkpoint_max: The maximum number of recent checkpoint files to keep. As new files are created, older files are deleted. If `None` or 0, all checkpoint files are kept. Defaults to 5 (that is, the 5 most recent checkpoint files are kept). If a saver is passed to the estimator, this argument will be ignored. keep_checkpoint_every_n_hours: Number of hours between each checkpoint to be saved. The default value of 10,000 hours effectively disables the feature. log_step_count_steps: The frequency, in number of global steps, that the global step and the loss will be logged during training. Also controls the frequency that the global steps / s will be logged (and written to summary) during training. train_distribute: An optional instance of `tf.distribute.Strategy`. If specified, then Estimator will distribute the user's model during training, according to the policy specified by that strategy. Setting `experimental_distribute.train_distribute` is preferred. device_fn: A callable invoked for every `Operation` that takes the `Operation` and returns the device string. If `None`, defaults to the device function returned by `tf.train.replica_device_setter` with round-robin strategy. protocol: An optional argument which specifies the protocol used when starting server. `None` means default to grpc. eval_distribute: An optional instance of `tf.distribute.Strategy`. If specified, then Estimator will distribute the user's model during evaluation, according to the policy specified by that strategy. Setting `experimental_distribute.eval_distribute` is preferred. experimental_distribute: An optional `tf.contrib.distribute.DistributeConfig` object specifying DistributionStrategy-related configuration. The `train_distribute` and `eval_distribute` can be passed as parameters to `RunConfig` or set in `experimental_distribute` but not both. experimental_max_worker_delay_secs: An optional integer specifying the maximum time a worker should wait before starting. By default, workers are started at staggered times, with each worker being delayed by up to 60 seconds. This is intended to reduce the risk of divergence, which can occur when many workers simultaneously update the weights of a randomly initialized model. Users who warm-start their models and train them for short durations (a few minutes or less) should consider reducing this default to improve training times. session_creation_timeout_secs: Max time workers should wait for a session to become available (on initialization or when recovering a session) with MonitoredTrainingSession. Defaults to 7200 seconds, but users may want to set a lower value to detect problems with variable / session (re)-initialization more quickly. checkpoint_save_graph_def: Whether to save the GraphDef and MetaGraphDef to `checkpoint_dir`. The GraphDef is saved after the session is created as `graph.pbtxt`. MetaGraphDefs are saved out for every checkpoint as `model.ckpt-*.meta`. Raises: ValueError: If both `save_checkpoints_steps` and `save_checkpoints_secs` are set. """ if (save_checkpoints_steps == _USE_DEFAULT and save_checkpoints_secs == _USE_DEFAULT): save_checkpoints_steps = None save_checkpoints_secs = 600 elif save_checkpoints_secs == _USE_DEFAULT: save_checkpoints_secs = None elif save_checkpoints_steps == _USE_DEFAULT: save_checkpoints_steps = None elif (save_checkpoints_steps is not None and save_checkpoints_secs is not None): raise ValueError(_SAVE_CKPT_ERR) self._verify_strategy_compatibility(train_distribute, eval_distribute) tf_config = json.loads(os.environ.get(_TF_CONFIG_ENV, '{}')) if tf_config: tf.compat.v1.logging.info('TF_CONFIG environment variable: %s', tf_config) model_dir = _get_model_dir(tf_config, path_to_str(model_dir)) RunConfig._replace( self, allowed_properties_list=_DEFAULT_REPLACEABLE_LIST, model_dir=model_dir, tf_random_seed=tf_random_seed, save_summary_steps=save_summary_steps, save_checkpoints_steps=save_checkpoints_steps, save_checkpoints_secs=save_checkpoints_secs, session_config=session_config, keep_checkpoint_max=keep_checkpoint_max, keep_checkpoint_every_n_hours=keep_checkpoint_every_n_hours, log_step_count_steps=log_step_count_steps, train_distribute=train_distribute, device_fn=device_fn, protocol=protocol, eval_distribute=eval_distribute, experimental_distribute=experimental_distribute, experimental_max_worker_delay_secs=experimental_max_worker_delay_secs, session_creation_timeout_secs=session_creation_timeout_secs, checkpoint_save_graph_def=checkpoint_save_graph_def) # TODO(frankchn,priyag): Eventually use distributed coordinator for TPUs. if ((train_distribute and not train_distribute.__class__.__name__.startswith('TPUStrategy')) or (eval_distribute and not eval_distribute.__class__.__name__.startswith('TPUStrategy')) or experimental_distribute): tf.compat.v1.logging.info( 'Initializing RunConfig with distribution strategies.') distribute_coordinator_training.init_run_config(self, tf_config) else: self._init_distributed_setting_from_environment_var(tf_config) self._maybe_overwrite_session_config_for_distributed_training() def _verify_strategy_compatibility(self, train_distribute, eval_distribute): if ((train_distribute is not None and train_distribute.__class__ == tf.compat.v2.distribute.experimental.ParameterServerStrategy) or (eval_distribute is not None and eval_distribute.__class__ == tf.compat.v2.distribute.experimental.ParameterServerStrategy)): raise ValueError('Please use `tf.compat.v1.distribute.experimental.Param' 'eterServerStrategy` for parameter server strategy with ' 'estimator.') def _maybe_overwrite_session_config_for_distributed_training(self): """Overwrites the session_config for distributed training. The default overwrite is optimized for between-graph training. Subclass should override this method if necessary. """ # Get session_config only for between-graph distributed mode (cluster_spec # is present). if not self._session_config and self._cluster_spec: RunConfig._replace( self, allowed_properties_list=_DEFAULT_REPLACEABLE_LIST, session_config=self._get_default_session_config_distributed()) def _get_default_session_config_distributed(self): """Returns None or tf.ConfigProto instance with default device_filters set. Device filters are set such that chief/master and worker communicates with only ps. session_config=None for evaluators or any other TaskType. """ rewrite_opts = rewriter_config_pb2.RewriterConfig( meta_optimizer_iterations=rewriter_config_pb2.RewriterConfig.ONE) graph_opts = tf.compat.v1.GraphOptions(rewrite_options=rewrite_opts) device_filters = None if self._task_type == TaskType.MASTER: device_filters = ['/job:ps', '/job:master'] elif self._task_type == TaskType.CHIEF: device_filters = ['/job:ps', '/job:chief'] elif self._task_type == TaskType.WORKER: device_filters = ['/job:ps', '/job:worker/task:%d' % self._task_id] elif self._task_type == TaskType.PS: device_filters = ['/job:ps', '/job:worker', '/job:chief', '/job:master'] else: # If the task_type is `EVALUATOR` or something other than the ones in # TaskType then don't set any device filters. return None return tf.compat.v1.ConfigProto( allow_soft_placement=True, graph_options=graph_opts, device_filters=device_filters) def _init_distributed_setting_from_environment_var(self, tf_config): """Initialize distributed properties based on `tf_config`.""" self._service = _validate_service(tf_config.get(_SERVICE_KEY)) self._cluster_spec = tf.train.ClusterSpec(tf_config.get(_CLUSTER_KEY, {})) task_env = tf_config.get(_TASK_ENV_KEY, {}) if self._cluster_spec and TaskType.MASTER in self._cluster_spec.jobs: return self._init_distributed_setting_from_environment_var_with_master( tf_config) if self._cluster_spec: # Distributed mode. self._task_type, self._task_id = _validate_task_type_and_task_id( self._cluster_spec, task_env, TaskType.CHIEF) self._evaluation_master = _get_eval_session_master( self._task_type, tf_config) if self._task_type != TaskType.EVALUATOR: self._master = _get_session_master(self._cluster_spec, self._task_type, self._task_id, tf_config) self._num_ps_replicas = _count_ps(self._cluster_spec) self._num_worker_replicas = _count_worker( self._cluster_spec, chief_task_type=TaskType.CHIEF) self._global_id_in_cluster = _get_global_id_in_cluster( self._cluster_spec, self._task_type, self._task_id, chief_task_type=TaskType.CHIEF) else: # Evaluator is not part of the training cluster. self._cluster_spec = tf.train.ClusterSpec({}) self._master = _LOCAL_MASTER self._num_ps_replicas = 0 self._num_worker_replicas = 0 self._global_id_in_cluster = None # undefined self._is_chief = self._task_type == TaskType.CHIEF else: # Local mode. self._task_type = task_env.get(_TASK_TYPE_KEY, TaskType.WORKER) self._task_id = int(task_env.get(_TASK_ID_KEY, 0)) self._global_id_in_cluster = 0 if self._task_type != TaskType.WORKER: raise ValueError( 'If "cluster" is not set in TF_CONFIG, task type must be WORKER.') if self._task_id != 0: raise ValueError( 'If "cluster" is not set in TF_CONFIG, task index must be 0.') self._master = tf_config.get(_SESSION_MASTER_KEY, _LOCAL_MASTER) self._evaluation_master = tf_config.get(_EVAL_SESSION_MASTER_KEY, _LOCAL_MASTER) self._is_chief = True self._num_ps_replicas = 0 self._num_worker_replicas = 1 def _init_distributed_setting_from_environment_var_with_master( self, tf_config): """Initialize distributed properties for legacy cluster with `master`.""" # There is no tech reason, why user cannot have chief and master in the same # cluster, but it is super confusing (which is really the chief?). So, block # this case. if TaskType.CHIEF in self._cluster_spec.jobs: raise ValueError('If `master` node exists in `cluster`, job ' '`chief` is not supported.') task_env = tf_config.get(_TASK_ENV_KEY, {}) self._task_type, self._task_id = _validate_task_type_and_task_id( self._cluster_spec, task_env, TaskType.MASTER) if self._task_type == TaskType.EVALUATOR: raise ValueError('If `master` node exists in `cluster`, task_type ' '`evaluator` is not supported.') self._global_id_in_cluster = _get_global_id_in_cluster( self._cluster_spec, self._task_type, self._task_id, chief_task_type=TaskType.MASTER) self._master = _get_session_master(self._cluster_spec, self._task_type, self._task_id, tf_config) self._evaluation_master = _get_eval_session_master(self._task_type, tf_config) self._num_ps_replicas = _count_ps(self._cluster_spec) self._num_worker_replicas = _count_worker( self._cluster_spec, chief_task_type=TaskType.MASTER) self._is_chief = self._task_type == TaskType.MASTER @property def cluster_spec(self): return self._cluster_spec @property def device_fn(self): """Returns the device_fn. If device_fn is not `None`, it overrides the default device function used in `Estimator`. Otherwise the default one is used. """ return self._device_fn @property def evaluation_master(self): return self._evaluation_master @property def is_chief(self): return self._is_chief @property def master(self): return self._master @property def num_ps_replicas(self): return self._num_ps_replicas @property def num_worker_replicas(self): return self._num_worker_replicas @property def task_id(self): return self._task_id @property def global_id_in_cluster(self): """The global id in the training cluster. All global ids in the training cluster are assigned from an increasing sequence of consecutive integers. The first id is 0. Note: Task id (the property field `task_id`) is tracking the index of the node among all nodes with the SAME task type. For example, given the cluster definition as follows: ``` cluster = {'chief': ['host0:2222'], 'ps': ['host1:2222', 'host2:2222'], 'worker': ['host3:2222', 'host4:2222', 'host5:2222']} ``` Nodes with task type `worker` can have id 0, 1, 2. Nodes with task type `ps` can have id, 0, 1. So, `task_id` is not unique, but the pair (`task_type`, `task_id`) can uniquely determine a node in the cluster. Global id, i.e., this field, is tracking the index of the node among ALL nodes in the cluster. It is uniquely assigned. For example, for the cluster spec given above, the global ids are assigned as: ``` task_type | task_id | global_id -------------------------------- chief | 0 | 0 worker | 0 | 1 worker | 1 | 2 worker | 2 | 3 ps | 0 | 4 ps | 1 | 5 ``` Returns: An integer id. """ return self._global_id_in_cluster @property def experimental_max_worker_delay_secs(self): return self._experimental_max_worker_delay_secs @property def task_type(self): return self._task_type @property def tf_random_seed(self): return self._tf_random_seed @property def save_summary_steps(self): return self._save_summary_steps @property def save_checkpoints_secs(self): return self._save_checkpoints_secs @property def session_config(self): return self._session_config @property def save_checkpoints_steps(self): return self._save_checkpoints_steps @property def checkpoint_save_graph_def(self): return self._checkpoint_save_graph_def @property def keep_checkpoint_max(self): return self._keep_checkpoint_max @property def session_creation_timeout_secs(self): return self._session_creation_timeout_secs @property def keep_checkpoint_every_n_hours(self): return self._keep_checkpoint_every_n_hours @property def log_step_count_steps(self): return self._log_step_count_steps @property def model_dir(self): return self._model_dir @property def service(self): """Returns the platform defined (in TF_CONFIG) service dict.""" return self._service @property def train_distribute(self): """Optional `tf.distribute.Strategy` for training.""" return self._train_distribute @property def eval_distribute(self): """Optional `tf.distribute.Strategy` for evaluation.""" return self._eval_distribute @property def protocol(self): """Returns the optional protocol value.""" return self._protocol def replace(self, **kwargs): """Returns a new instance of `RunConfig` replacing specified properties. Only the properties in the following list are allowed to be replaced: - `model_dir`, - `tf_random_seed`, - `save_summary_steps`, - `save_checkpoints_steps`, - `save_checkpoints_secs`, - `session_config`, - `keep_checkpoint_max`, - `keep_checkpoint_every_n_hours`, - `log_step_count_steps`, - `train_distribute`, - `device_fn`, - `protocol`. - `eval_distribute`, - `experimental_distribute`, - `experimental_max_worker_delay_secs`, In addition, either `save_checkpoints_steps` or `save_checkpoints_secs` can be set (should not be both). Args: **kwargs: keyword named properties with new values. Raises: ValueError: If any property name in `kwargs` does not exist or is not allowed to be replaced, or both `save_checkpoints_steps` and `save_checkpoints_secs` are set. Returns: a new instance of `RunConfig`. """ return RunConfig._replace( copy.deepcopy(self), allowed_properties_list=_DEFAULT_REPLACEABLE_LIST, **kwargs) @staticmethod def _replace(config, allowed_properties_list=None, **kwargs): """See `replace`. N.B.: This implementation assumes that for key named "foo", the underlying property the RunConfig holds is "_foo" (with one leading underscore). Args: config: The RunConfig to replace the values of. allowed_properties_list: The property name list allowed to be replaced. **kwargs: keyword named properties with new values. Raises: ValueError: If any property name in `kwargs` does not exist or is not allowed to be replaced, or both `save_checkpoints_steps` and `save_checkpoints_secs` are set. Returns: a new instance of `RunConfig`. """ allowed_properties_list = allowed_properties_list or [] for key, new_value in six.iteritems(kwargs): if key in allowed_properties_list: setattr(config, '_' + key, new_value) continue raise ValueError( 'Replacing {} is not supported. Allowed properties are {}.'.format( key, allowed_properties_list)) _validate_save_ckpt_with_replaced_keys(config, kwargs.keys()) _validate_properties(config) return config def _get_model_dir(tf_config, model_dir): """Returns `model_dir` based user provided `tf_config` or `model_dir`.""" # pylint: disable=g-explicit-bool-comparison # Empty string is treated as False in Python condition check, which triggers # some confusing error messages. For example, 'a or b' returns None if a is '' # and b is None. `None` is allowed for model_dir but '' is not allowed. Here, # explicitly check empty string to provide clear error message. if model_dir == '': raise ValueError('model_dir should be non-empty.') model_dir_in_tf_config = tf_config.get('model_dir') if model_dir_in_tf_config == '': raise ValueError('model_dir in TF_CONFIG should be non-empty.') if model_dir_in_tf_config: if model_dir and model_dir_in_tf_config != model_dir: raise ValueError( '`model_dir` provided in RunConfig construct, if set, ' 'must have the same value as the model_dir in TF_CONFIG. ' 'model_dir: {}\nTF_CONFIG["model_dir"]: {}.\n'.format( model_dir, model_dir_in_tf_config)) tf.compat.v1.logging.info('Using model_dir in TF_CONFIG: %s', model_dir_in_tf_config) return model_dir or model_dir_in_tf_config def path_to_str(path): """Returns the file system path representation of a `PathLike` object, else as it is. Args: path: An object that can be converted to path representation. Returns: A `str` object. """ if hasattr(path, '__fspath__'): path = tf.compat.as_str_any(path.__fspath__()) return path