ParameterServerStrategy worker removing code change

[wojciechp6]

Please make sure that this is a feature request. As per our GitHub Policy, we only address code/doc bugs, performance issues, feature requests and build/installation issues on GitHub. tag:feature_template

System information

• TensorFlow version (you are using): 2.4.1
• Are you willing to contribute it (Yes/No): No

Hi, I'm working on dynamic worker adding and removing for ParameterServerStrategy. I noticed that one change in Worker class is necessary to remove worker while training.
The change is here:
cluster_coordinator.py

class Worker:
  [...]
  def _process_queue(self):
    """Function running in a worker thread to process closure queues."""
    self._maybe_delay()
    while self._should_worker_thread_run:
      closure = self._cluster._closure_queue.get()  # pylint: disable=protected-access
      if not self._should_worker_thread_run:
        closure.mark_cancelled()
        return
      if closure is None:
        return
      self._process_closure(closure)
      # To properly stop the worker and preemption threads, it is important that
      # `ClusterCoordinator` object is not held onto so its `__del__` can be
      # called. By removing the reference to the `closure` that has already been
      # processed, we ensure that the `closure` object is released, while
      # getting the next `closure` at above `self._cluster._closure_queue.get()`
      # call.
      del closure
  [...]

Could someone commit it? I'm not familiar with tensorflow contribute requirements, so I think somebody else could do this faster. Thanks in advance.

[yuefengz]

Looks like you want to add the if block containing closure.mark_cancelled()? Is this the only change? Do you want to send a PR? I am happy to change it as well.

[wojciechp6]

Hi! Yes, the if block is the only change. I'll be happy if you could do it for me. The contribution polices are too complex to merge such small change. Greetings!