def _update_last_index_token(self, new_index_token: str) -> None: """ This function updates the last_index_token in the following scenarios: 1. last_complete_step >= last_index_token_step : this means that the token isn't pointing to the latest completed step 2. number of steps available ( complete or incomplete ) - (last_completed_step+1) > window_size_limit: we maintain a window to stop querying for older steps that have not completed. if the total number of steps, we are querying for completion is greater than our window_size_limit we update the last_index_token and last_complete_step by (window_size_limit // 2) :param new_index_token: :return:None """ if self.last_index_token is None: last_index_token_step = 0 else: last_index_token_step = IndexFileLocationUtils.parse_step_from_index_file_name( self.last_index_token) # Case 1: This case is not satisfied when all workers in a # distributed training job have not written a step if self.last_complete_step >= last_index_token_step: prefix = IndexFileLocationUtils.get_prefix_from_index_file( new_index_token) # sort lexicographically and select the last worker last_worker = sorted(list(self.worker_set))[-1] # below converts worker_name to serialized workerName # if it's a tf device, else no effect last_worker_serialized = serialize_tf_device(last_worker) self.last_index_token = IndexFileLocationUtils.get_index_key_for_step( prefix, self.last_complete_step, last_worker_serialized) self.logger.debug( f"Updated last index token to:{self.last_index_token}") # Case 2: This case is satisfied if the number of incomplete steps # is greater than the INCOMPLETE_STEP_WAIT_WINDOW available_step = self._global_to_mode.keys() if (len(available_step) - (self.last_complete_step + 1) > self._incomplete_wait_for_step_window): prefix = IndexFileLocationUtils.get_prefix_from_index_file( new_index_token) last_worker = sorted(list(self.worker_set))[-1] # below converts worker_name to serialized workerName # if it's a tf device, else no effect last_worker_serialized = serialize_tf_device(last_worker) self.last_index_token = IndexFileLocationUtils.get_index_key_for_step( prefix, self.last_complete_step + (self._incomplete_wait_for_step_window // 2), last_worker_serialized, ) self.last_complete_step = IndexFileLocationUtils.parse_step_from_index_file_name( self.last_index_token) self.logger.info( f"Waiting for: {len(available_step) - (self.last_complete_step + 1)} Steps. \n" f"INCOMPLETE_STEP_WAIT_WINDOW: {self._incomplete_wait_for_step_window}. \n" f"Marking the last {self._incomplete_wait_for_step_window // 2} incomplete steps as complete" f"Updating last_index_token to: {self.last_index_token}. \n" f"Updating last_complete_step to: {self.last_complete_step}. ")
def test_get_prefix_from_index_file(): local_index_filepath = "/opt/ml/testing/run_1/index/000000000/000000000000_worker_0.json" prefix = IndexFileLocationUtils.get_prefix_from_index_file(local_index_filepath) assert prefix == "/opt/ml/testing/run_1" s3_index_filepath = ( "s3://bucket-that-does-not-exist/run_1/index/000000000/000000000000_worker_0.json" ) prefix = IndexFileLocationUtils.get_prefix_from_index_file(s3_index_filepath) assert prefix == "s3://bucket-that-does-not-exist/run_1"