data_preprocessing.py 26.5 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
# Copyright 2018 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.
# ==============================================================================
"""Preprocess dataset and construct any necessary artifacts."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import atexit
import contextlib
import gc
24
import hashlib
25
26
27
28
29
import multiprocessing
import json
import os
import pickle
import signal
30
import socket
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
import subprocess
import time
import timeit
import typing

# pylint: disable=wrong-import-order
from absl import app as absl_app
from absl import flags
import numpy as np
import pandas as pd
import six
import tensorflow as tf
# pylint: enable=wrong-import-order

from official.datasets import movielens
from official.recommendation import constants as rconst
from official.recommendation import stat_utils
48
from official.recommendation import popen_helper
49
50


51
52
53
54
55
56
57
58
59
DATASET_TO_NUM_USERS_AND_ITEMS = {
    "ml-1m": (6040, 3706),
    "ml-20m": (138493, 26744)
}


# Number of batches to run per epoch when using synthetic data. At high batch
# sizes, we run for more batches than with real data, which is good since
# running more batches reduces noise when measuring the average batches/second.
60
SYNTHETIC_BATCHES_PER_EPOCH = 2000
61
62


63
64
65
66
class NCFDataset(object):
  """Container for training and testing data."""

  def __init__(self, user_map, item_map, num_data_readers, cache_paths,
67
               num_train_positives, deterministic=False):
68
    # type: (dict, dict, int, rconst.Paths, int, bool) -> None
69
70
71
72
73
74
75
    """Assign key values for recommendation dataset.

    Args:
      user_map: Dict mapping raw user ids to regularized ids.
      item_map: Dict mapping raw item ids to regularized ids.
      num_data_readers: The number of reader Datasets used during training.
      cache_paths: Object containing locations for various cache files.
76
77
      num_train_positives: The number of positive training examples in the
        dataset.
78
79
      deterministic: Operations should use deterministic, order preserving
        methods, even at the cost of performance.
80
81
82
83
84
85
86
87
88
    """

    self.user_map = {int(k): int(v) for k, v in user_map.items()}
    self.item_map = {int(k): int(v) for k, v in item_map.items()}
    self.num_users = len(user_map)
    self.num_items = len(item_map)
    self.num_data_readers = num_data_readers
    self.cache_paths = cache_paths
    self.num_train_positives = num_train_positives
89
    self.deterministic = deterministic
90
91


92
93
def _filter_index_sort(raw_rating_path, match_mlperf):
  # type: (str, bool) -> (pd.DataFrame, dict, dict)
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
  """Read in data CSV, and output structured data.

  This function reads in the raw CSV of positive items, and performs three
  preprocessing transformations:

  1)  Filter out all users who have not rated at least a certain number
      of items. (Typically 20 items)

  2)  Zero index the users and items such that the largest user_id is
      `num_users - 1` and the largest item_id is `num_items - 1`

  3)  Sort the dataframe by user_id, with timestamp as a secondary sort key.
      This allows the dataframe to be sliced by user in-place, and for the last
      item to be selected simply by calling the `-1` index of a user's slice.

  While all of these transformations are performed by Pandas (and are therefore
  single-threaded), they only take ~2 minutes, and the overhead to apply a
  MapReduce pattern to parallel process the dataset adds significant complexity
  for no computational gain. For a larger dataset parallelizing this
  preprocessing could yield speedups. (Also, this preprocessing step is only
  performed once for an entire run.

  Args:
    raw_rating_path: The path to the CSV which contains the raw dataset.
118
119
    match_mlperf: If True, change the sorting algorithm to match the MLPerf
      reference implementation.
120
121

  Returns:
Reed's avatar
Reed committed
122
123
124
    A filtered, zero-index remapped, sorted dataframe, a dict mapping raw user
    IDs to regularized user IDs, and a dict mapping raw item IDs to regularized
    item IDs.
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
  """
  with tf.gfile.Open(raw_rating_path) as f:
    df = pd.read_csv(f)

  # Get the info of users who have more than 20 ratings on items
  grouped = df.groupby(movielens.USER_COLUMN)
  df = grouped.filter(
      lambda x: len(x) >= rconst.MIN_NUM_RATINGS) # type: pd.DataFrame

  original_users = df[movielens.USER_COLUMN].unique()
  original_items = df[movielens.ITEM_COLUMN].unique()

  # Map the ids of user and item to 0 based index for following processing
  tf.logging.info("Generating user_map and item_map...")
  user_map = {user: index for index, user in enumerate(original_users)}
  item_map = {item: index for index, item in enumerate(original_items)}

  df[movielens.USER_COLUMN] = df[movielens.USER_COLUMN].apply(
      lambda user: user_map[user])
  df[movielens.ITEM_COLUMN] = df[movielens.ITEM_COLUMN].apply(
      lambda item: item_map[item])

  num_users = len(original_users)
  num_items = len(original_items)

  assert num_users <= np.iinfo(np.int32).max
  assert num_items <= np.iinfo(np.uint16).max
  assert df[movielens.USER_COLUMN].max() == num_users - 1
  assert df[movielens.ITEM_COLUMN].max() == num_items - 1

  # This sort is used to shard the dataframe by user, and later to select
  # the last item for a user to be used in validation.
  tf.logging.info("Sorting by user, timestamp...")
158
159
160
161
162
163
164
165
166
167
168
169

  if match_mlperf:
    # This sort is equivalent to the non-MLPerf sort, except that the order of
    # items with the same user and timestamp are sometimes different. For some
    # reason, this sort results in a better hit-rate during evaluation, matching
    # the performance of the MLPerf reference implementation.
    df.sort_values(by=movielens.TIMESTAMP_COLUMN, inplace=True)
    df.sort_values([movielens.USER_COLUMN, movielens.TIMESTAMP_COLUMN],
                   inplace=True, kind="mergesort")
  else:
    df.sort_values([movielens.USER_COLUMN, movielens.TIMESTAMP_COLUMN],
                   inplace=True)
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187

  df = df.reset_index()  # The dataframe does not reconstruct indicies in the
  # sort or filter steps.

  return df, user_map, item_map


def _train_eval_map_fn(args):
  """Split training and testing data and generate testing negatives.

  This function is called as part of a multiprocessing map. The principle
  input is a shard, which contains a sorted array of users and corresponding
  items for each user, where items have already been sorted in ascending order
  by timestamp. (Timestamp is not passed to avoid the serialization cost of
  sending it to the map function.)

  For each user, all but the last item is written into a pickle file which the
  training data producer can consume on as needed. The last item for a user
188
189
  is a validation point; it is written under a separate key and will be used
  later to generate the evaluation data.
190
191
192
193
194
195
196

  Args:
    shard: A dict containing the user and item arrays.
    shard_id: The id of the shard provided. This is used to number the training
      shard pickle files.
    num_items: The cardinality of the item set, which determines the set from
      which validation negatives should be drawn.
Reed's avatar
Reed committed
197
198
    cache_paths: rconst.Paths object containing locations for various cache
      files.
199
200
201

  """

202
  shard, shard_id, num_items, cache_paths = args
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231

  users = shard[movielens.USER_COLUMN]
  items = shard[movielens.ITEM_COLUMN]

  # This produces index boundaries which can be used to slice by user.
  delta = users[1:] - users[:-1]
  boundaries = ([0] + (np.argwhere(delta)[:, 0] + 1).tolist() +
                [users.shape[0]])

  train_blocks = []
  test_positives = []
  for i in range(len(boundaries) - 1):
    # This is simply a vector of repeated values such that the shard could be
    # represented compactly with a tuple of tuples:
    #   ((user_id, items), (user_id, items), ...)
    # rather than:
    #   user_id_vector, item_id_vector
    # However the additional nested structure significantly increases the
    # serialization and deserialization cost such that it is not worthwhile.
    block_user = users[boundaries[i]:boundaries[i+1]]
    assert len(set(block_user)) == 1

    block_items = items[boundaries[i]:boundaries[i+1]]
    train_blocks.append((block_user[:-1], block_items[:-1]))
    test_positives.append((block_user[0], block_items[-1]))

  train_users = np.concatenate([i[0] for i in train_blocks])
  train_items = np.concatenate([i[1] for i in train_blocks])

232
233
234
235
236
  test_pos_users = np.array([i[0] for i in test_positives],
                            dtype=train_users.dtype)
  test_pos_items = np.array([i[1] for i in test_positives],
                            dtype=train_items.dtype)

237
238
239
240
241
  train_shard_fpath = cache_paths.train_shard_template.format(
      str(shard_id).zfill(5))

  with tf.gfile.Open(train_shard_fpath, "wb") as f:
    pickle.dump({
242
243
244
245
246
247
248
249
        rconst.TRAIN_KEY: {
            movielens.USER_COLUMN: train_users,
            movielens.ITEM_COLUMN: train_items,
        },
        rconst.EVAL_KEY: {
            movielens.USER_COLUMN: test_pos_users,
            movielens.ITEM_COLUMN: test_pos_items,
        }
250
251
252
    }, f)


253
254
255
def generate_train_eval_data(df, approx_num_shards, num_items, cache_paths,
                             match_mlperf):
  # type: (pd.DataFrame, int, int, rconst.Paths, bool) -> None
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
  """Construct training and evaluation datasets.

  This function manages dataset construction and validation that the
  transformations have produced correct results. The particular logic of
  transforming the data is performed in _train_eval_map_fn().

  Args:
    df: The dataframe containing the entire dataset. It is essential that this
      dataframe be produced by _filter_index_sort(), as subsequent
      transformations rely on `df` having particular structure.
    approx_num_shards: The approximate number of similarly sized shards to
      construct from `df`. The MovieLens has severe imbalances where some users
      have interacted with many items; this is common among datasets involving
      user data. Rather than attempt to aggressively balance shard size, this
      function simply allows shards to "overflow" which can produce a number of
      shards which is less than `approx_num_shards`. This small degree of
      imbalance does not impact performance; however it does mean that one
      should not expect approx_num_shards to be the ACTUAL number of shards.
    num_items: The cardinality of the item set.
Reed's avatar
Reed committed
275
276
    cache_paths: rconst.Paths object containing locations for various cache
      files.
277
278
    match_mlperf: If True, sample eval negative with replacements, which the
      MLPerf reference implementation does.
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
  """

  num_rows = len(df)
  approximate_partitions = np.linspace(
      0, num_rows, approx_num_shards + 1).astype("int")
  start_ind, end_ind = 0, 0
  shards = []

  for i in range(1, approx_num_shards + 1):
    end_ind = approximate_partitions[i]
    while (end_ind < num_rows and df[movielens.USER_COLUMN][end_ind - 1] ==
           df[movielens.USER_COLUMN][end_ind]):
      end_ind += 1

    if end_ind <= start_ind:
      continue  # imbalance from prior shard.

    df_shard = df[start_ind:end_ind]
    user_shard = df_shard[movielens.USER_COLUMN].values.astype(np.int32)
    item_shard = df_shard[movielens.ITEM_COLUMN].values.astype(np.uint16)

    shards.append({
        movielens.USER_COLUMN: user_shard,
        movielens.ITEM_COLUMN: item_shard,
    })

    start_ind = end_ind
  assert end_ind == num_rows
  approx_num_shards = len(shards)

  tf.logging.info("Splitting train and test data and generating {} test "
                  "negatives per user...".format(rconst.NUM_EVAL_NEGATIVES))
  tf.gfile.MakeDirs(cache_paths.train_shard_subdir)

313
  map_args = [(shards[i], i, num_items, cache_paths)
314
              for i in range(approx_num_shards)]
315

316
317
  with popen_helper.get_pool(multiprocessing.cpu_count()) as pool:
    pool.map(_train_eval_map_fn, map_args)  # pylint: disable=no-member
318
319


320
def construct_cache(dataset, data_dir, num_data_readers, match_mlperf,
321
                    deterministic, cache_id=None):
322
  # type: (str, str, int, bool, bool, typing.Optional[int]) -> NCFDataset
323
324
325
326
327
328
329
  """Load and digest data CSV into a usable form.

  Args:
    dataset: The name of the dataset to be used.
    data_dir: The root directory of the dataset.
    num_data_readers: The number of parallel processes which will request
      data during training.
330
331
    match_mlperf: If True, change the behavior of the cache construction to
      match the MLPerf reference implementation.
332
333
    deterministic: Try to enforce repeatable behavior, even at the cost of
      performance.
334
  """
335
  cache_paths = rconst.Paths(data_dir=data_dir, cache_id=cache_id)
336
337
  num_data_readers = (num_data_readers or int(multiprocessing.cpu_count() / 2)
                      or 1)
338
339
340
341
342
343
344
345
346
347
348
349
  approx_num_shards = int(movielens.NUM_RATINGS[dataset]
                          // rconst.APPROX_PTS_PER_TRAIN_SHARD) or 1

  st = timeit.default_timer()
  cache_root = os.path.join(data_dir, cache_paths.cache_root)
  if tf.gfile.Exists(cache_root):
    raise ValueError("{} unexpectedly already exists."
                     .format(cache_paths.cache_root))
  tf.logging.info("Creating cache directory. This should be deleted on exit.")
  tf.gfile.MakeDirs(cache_paths.cache_root)

  raw_rating_path = os.path.join(data_dir, dataset, movielens.RATINGS_FILE)
350
  df, user_map, item_map = _filter_index_sort(raw_rating_path, match_mlperf)
351
352
353
354
355
356
357
358
  num_users, num_items = DATASET_TO_NUM_USERS_AND_ITEMS[dataset]

  if num_users != len(user_map):
    raise ValueError("Expected to find {} users, but found {}".format(
        num_users, len(user_map)))
  if num_items != len(item_map):
    raise ValueError("Expected to find {} items, but found {}".format(
        num_items, len(item_map)))
359
360

  generate_train_eval_data(df=df, approx_num_shards=approx_num_shards,
361
362
                           num_items=len(item_map), cache_paths=cache_paths,
                           match_mlperf=match_mlperf)
363
364
365
366
367
  del approx_num_shards  # value may have changed.

  ncf_dataset = NCFDataset(user_map=user_map, item_map=item_map,
                           num_data_readers=num_data_readers,
                           cache_paths=cache_paths,
368
369
                           num_train_positives=len(df) - len(user_map),
                           deterministic=deterministic)
370
371
372
373
374
375
376
377
378
379
380
381
382

  run_time = timeit.default_timer() - st
  tf.logging.info("Cache construction complete. Time: {:.1f} sec."
                  .format(run_time))

  return ncf_dataset


def _shutdown(proc):
  # type: (subprocess.Popen) -> None
  """Convenience function to cleanly shut down async generation process."""

  tf.logging.info("Shutting down train data creation subprocess.")
383
  try:
384
385
386
387
388
    try:
      proc.send_signal(signal.SIGINT)
      time.sleep(5)
      if proc.returncode is not None:
        return  # SIGINT was handled successfully within 5 seconds
389

390
391
    except socket.error:
      pass
392

393
394
395
396
397
    # Otherwise another second of grace period and then force kill the process.
    time.sleep(1)
    proc.terminate()
  except:  # pylint: disable=broad-except
    tf.logging.error("Data generation subprocess could not be killed.")
398
399
400


def instantiate_pipeline(dataset, data_dir, batch_size, eval_batch_size,
401
                         num_data_readers=None, num_neg=4, epochs_per_cycle=1,
402
                         match_mlperf=False, deterministic=False,
403
                         use_subprocess=True, cache_id=None):
404
  # type: (...) -> (NCFDataset, typing.Callable)
405
406
407
408
  """Preprocess data and start negative generation subprocess."""

  tf.logging.info("Beginning data preprocessing.")
  ncf_dataset = construct_cache(dataset=dataset, data_dir=data_dir,
409
                                num_data_readers=num_data_readers,
410
                                match_mlperf=match_mlperf,
411
412
                                deterministic=deterministic,
                                cache_id=cache_id)
413
414
  # By limiting the number of workers we guarantee that the worker
  # pool underlying the training generation doesn't starve other processes.
415
  num_workers = int(multiprocessing.cpu_count() * 0.75) or 1
416

417
418
419
420
421
422
  flags_ = {
      "data_dir": data_dir,
      "cache_id": ncf_dataset.cache_paths.cache_id,
      "num_neg": num_neg,
      "num_train_positives": ncf_dataset.num_train_positives,
      "num_items": ncf_dataset.num_items,
423
      "num_users": ncf_dataset.num_users,
424
425
426
427
428
429
430
      "num_readers": ncf_dataset.num_data_readers,
      "epochs_per_cycle": epochs_per_cycle,
      "train_batch_size": batch_size,
      "eval_batch_size": eval_batch_size,
      "num_workers": num_workers,
      "redirect_logs": use_subprocess,
      "use_tf_logging": not use_subprocess,
431
      "ml_perf": match_mlperf,
432
  }
433

434
435
  if ncf_dataset.deterministic:
    flags_["seed"] = stat_utils.random_int32()
436
  tf.gfile.MakeDirs(data_dir)
437
438
439
  # We write to a temp file then atomically rename it to the final file,
  # because writing directly to the final file can cause the data generation
  # async process to read a partially written JSON file.
440
441
  flagfile_temp = os.path.join(ncf_dataset.cache_paths.cache_root,
                               rconst.FLAGFILE_TEMP)
442
443
444
445
446
  tf.logging.info("Preparing flagfile for async data generation in {} ..."
                  .format(flagfile_temp))
  with tf.gfile.Open(flagfile_temp, "w") as f:
    for k, v in six.iteritems(flags_):
      f.write("--{}={}\n".format(k, v))
447
  flagfile = os.path.join(ncf_dataset.cache_paths.cache_root, rconst.FLAGFILE)
448
449
450
451
452
  tf.gfile.Rename(flagfile_temp, flagfile)
  tf.logging.info(
      "Wrote flagfile for async data generation in {}."
      .format(flagfile))

453
454
455
456
457
458
459
460
  if use_subprocess:
    tf.logging.info("Creating training file subprocess.")
    subproc_env = os.environ.copy()
    # The subprocess uses TensorFlow for tf.gfile, but it does not need GPU
    # resources and by default will try to allocate GPU memory. This would cause
    # contention with the main training process.
    subproc_env["CUDA_VISIBLE_DEVICES"] = ""
    subproc_args = popen_helper.INVOCATION + [
461
462
        "--data_dir", data_dir,
        "--cache_id", str(ncf_dataset.cache_paths.cache_id)]
463
464
465
466
    tf.logging.info(
        "Generation subprocess command: {}".format(" ".join(subproc_args)))
    proc = subprocess.Popen(args=subproc_args, shell=False, env=subproc_env)

467
468
469
470
471
472
473
  cleanup_called = {"finished": False}
  @atexit.register
  def cleanup():
    """Remove files and subprocess from data generation."""
    if cleanup_called["finished"]:
      return

474
475
476
    if use_subprocess:
      _shutdown(proc)

477
478
479
480
481
482
    try:
      tf.gfile.DeleteRecursively(ncf_dataset.cache_paths.cache_root)
    except tf.errors.NotFoundError:
      pass

    cleanup_called["finished"] = True
483

484
  for _ in range(300):
485
486
487
488
489
490
491
    if tf.gfile.Exists(ncf_dataset.cache_paths.subproc_alive):
      break
    time.sleep(1)  # allow `alive` file to be written
  if not tf.gfile.Exists(ncf_dataset.cache_paths.subproc_alive):
    raise ValueError("Generation subprocess did not start correctly. Data will "
                     "not be available; exiting to avoid waiting forever.")

492
  return ncf_dataset, cleanup
493
494
495
496
497
498
499
500
501
502


def make_deserialize(params, batch_size, training=False):
  """Construct deserialize function for training and eval fns."""
  feature_map = {
      movielens.USER_COLUMN: tf.FixedLenFeature([], dtype=tf.string),
      movielens.ITEM_COLUMN: tf.FixedLenFeature([], dtype=tf.string),
  }
  if training:
    feature_map["labels"] = tf.FixedLenFeature([], dtype=tf.string)
503
504
  else:
    feature_map[rconst.DUPLICATE_MASK] = tf.FixedLenFeature([], dtype=tf.string)
505
506
507
508
509
510
511
512
513
514
515
516
517

  def deserialize(examples_serialized):
    """Called by Dataset.map() to convert batches of records to tensors."""
    features = tf.parse_single_example(examples_serialized, feature_map)
    users = tf.reshape(tf.decode_raw(
        features[movielens.USER_COLUMN], tf.int32), (batch_size,))
    items = tf.reshape(tf.decode_raw(
        features[movielens.ITEM_COLUMN], tf.uint16), (batch_size,))

    if params["use_tpu"]:
      items = tf.cast(items, tf.int32)  # TPU doesn't allow uint16 infeed.

    if not training:
518
519
      dupe_mask = tf.reshape(tf.cast(tf.decode_raw(
          features[rconst.DUPLICATE_MASK], tf.int8), tf.bool), (batch_size,))
520
521
522
      return {
          movielens.USER_COLUMN: users,
          movielens.ITEM_COLUMN: items,
523
          rconst.DUPLICATE_MASK: dupe_mask,
524
525
526
527
      }

    labels = tf.reshape(tf.cast(tf.decode_raw(
        features["labels"], tf.int8), tf.bool), (batch_size,))
528

529
530
531
532
533
534
535
    return {
        movielens.USER_COLUMN: users,
        movielens.ITEM_COLUMN: items,
    }, labels
  return deserialize


536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
def hash_pipeline(dataset, deterministic):
  # type: (tf.data.Dataset, bool) -> None
  """Utility function for detecting non-determinism in the data pipeline.

  Args:
    dataset: a tf.data.Dataset generated by the input_fn
    deterministic: Does the input_fn expect the dataset to be deterministic.
      (i.e. fixed seed, sloppy=False, etc.)
  """
  if not deterministic:
    tf.logging.warning("Data pipeline is not marked as deterministic. Hash "
                       "values are not expected to be meaningful.")

  batch = dataset.make_one_shot_iterator().get_next()
  md5 = hashlib.md5()
  count = 0
  first_batch_hash = b""
  with tf.Session() as sess:
    while True:
      try:
        result = sess.run(batch)
        if isinstance(result, tuple):
          result = result[0]  # only hash features
      except tf.errors.OutOfRangeError:
        break

      count += 1
      md5.update(memoryview(result[movielens.USER_COLUMN]).tobytes())
      md5.update(memoryview(result[movielens.ITEM_COLUMN]).tobytes())
      if count == 1:
        first_batch_hash = md5.hexdigest()
  overall_hash = md5.hexdigest()
  tf.logging.info("Batch count: {}".format(count))
  tf.logging.info("  [pipeline_hash] First batch hash: {}".format(
      first_batch_hash))
  tf.logging.info("  [pipeline_hash] All batches hash: {}".format(overall_hash))


574
575
def make_input_fn(ncf_dataset, is_training):
  # type: (typing.Optional[NCFDataset], bool) -> (typing.Callable, str, int)
576
577
  """Construct training input_fn for the current epoch."""

578
  if ncf_dataset is None:
579
    return make_synthetic_input_fn(is_training)
580

581
  if not tf.gfile.Exists(ncf_dataset.cache_paths.subproc_alive):
582
583
584
585
    # The generation subprocess must have been alive at some point, because we
    # earlier checked that the subproc_alive file existed.
    raise ValueError("Generation subprocess unexpectedly died. Data will not "
                     "be available; exiting to avoid waiting forever.")
586

587
588
589
590
591
  if is_training:
    train_epoch_dir = ncf_dataset.cache_paths.train_epoch_dir
    while not tf.gfile.Exists(train_epoch_dir):
      tf.logging.info("Waiting for {} to exist.".format(train_epoch_dir))
      time.sleep(1)
592
593

    train_data_dirs = tf.gfile.ListDirectory(train_epoch_dir)
594
595
596
597
598
599
600
601
602
603
604
    while not train_data_dirs:
      tf.logging.info("Waiting for data folder to be created.")
      time.sleep(1)
      train_data_dirs = tf.gfile.ListDirectory(train_epoch_dir)
    train_data_dirs.sort()  # names are zfilled so that
                            # lexicographic sort == numeric sort
    record_dir = os.path.join(train_epoch_dir, train_data_dirs[0])
    template = rconst.TRAIN_RECORD_TEMPLATE
  else:
    record_dir = ncf_dataset.cache_paths.eval_data_subdir
    template = rconst.EVAL_RECORD_TEMPLATE
605
606
607
608
609
610
611
612
613

  ready_file = os.path.join(record_dir, rconst.READY_FILE)
  while not tf.gfile.Exists(ready_file):
    tf.logging.info("Waiting for records in {} to be ready".format(record_dir))
    time.sleep(1)

  with tf.gfile.Open(ready_file, "r") as f:
    epoch_metadata = json.load(f)

614
615
  # This value is used to check that the batch count from the subprocess matches
  # the batch count expected by the main thread.
616
617
618
619
  batch_count = epoch_metadata["batch_count"]

  def input_fn(params):
    """Generated input_fn for the given epoch."""
620
621
622
623
624
625
    if is_training:
      batch_size = params["batch_size"]
    else:
      # Estimator has "eval_batch_size" included in the params, but TPUEstimator
      # populates "batch_size" to the appropriate value.
      batch_size = params.get("eval_batch_size") or params["batch_size"]
626
627
628
629
630
631
632
633
634

    if epoch_metadata["batch_size"] != batch_size:
      raise ValueError(
          "Records were constructed with batch size {}, but input_fn was given "
          "a batch size of {}. This will result in a deserialization error in "
          "tf.parse_single_example."
          .format(epoch_metadata["batch_size"], batch_size))

    record_files = tf.data.Dataset.list_files(
635
        os.path.join(record_dir, template.format("*")), shuffle=False)
636
637
638
639
640

    interleave = tf.contrib.data.parallel_interleave(
        tf.data.TFRecordDataset,
        cycle_length=4,
        block_length=100000,
641
        sloppy=not ncf_dataset.deterministic,
642
643
644
        prefetch_input_elements=4,
    )

645
    deserialize = make_deserialize(params, batch_size, is_training)
646
647
    dataset = record_files.apply(interleave)
    dataset = dataset.map(deserialize, num_parallel_calls=4)
648
649
650
651
652
653
    dataset = dataset.prefetch(32)

    if params.get("hash_pipeline"):
      hash_pipeline(dataset, ncf_dataset.deterministic)

    return dataset
654
655
656
657

  return input_fn, record_dir, batch_count


658
def make_synthetic_input_fn(is_training):
659
660
661
  """Construct training input_fn that uses synthetic data."""
  def input_fn(params):
    """Generated input_fn for the given epoch."""
662
663
    batch_size = (params["batch_size"] if is_training else
                  params["eval_batch_size"] or params["batch_size"])
664
665
666
667
668
669
670
    num_users = params["num_users"]
    num_items = params["num_items"]

    users = tf.random_uniform([batch_size], dtype=tf.int32, minval=0,
                              maxval=num_users)
    items = tf.random_uniform([batch_size], dtype=tf.int32, minval=0,
                              maxval=num_items)
671

672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
    if is_training:
      labels = tf.random_uniform([batch_size], dtype=tf.int32, minval=0,
                                 maxval=2)
      data = {
          movielens.USER_COLUMN: users,
          movielens.ITEM_COLUMN: items,
      }, labels
    else:
      dupe_mask = tf.cast(tf.random_uniform([batch_size], dtype=tf.int32,
                                            minval=0, maxval=2), tf.bool)
      data = {
          movielens.USER_COLUMN: users,
          movielens.ITEM_COLUMN: items,
          rconst.DUPLICATE_MASK: dupe_mask,
      }
687
688

    dataset = tf.data.Dataset.from_tensors(data).repeat(
689
690
        SYNTHETIC_BATCHES_PER_EPOCH)
    dataset = dataset.prefetch(32)
691
692
    return dataset

693
  return input_fn, None, SYNTHETIC_BATCHES_PER_EPOCH