def _dist_search_fn(temp_dir, study_id, tuner_id):
    """Multi-process safe tuner instantiation and tuner.search()."""

    # Jitter instantiation so as to avoid contention on study
    # creation and dataset download.
    time.sleep(int(tuner_id[5:]))  # tuner_id is formatted as 'tuner%d'

    # Dataset must be loaded independently in sub-process.
    (x, y), (val_x, val_y) = _load_data()

    tuner = cloudtuner.CloudTuner(_build_model,
                                  project_id=_PROJECT_ID,
                                  region=_REGION,
                                  objective='acc',
                                  hyperparameters=_HPS,
                                  max_trials=5,
                                  study_id=study_id,
                                  directory=os.path.join(
                                      temp_dir, study_id, tuner_id))
    tuner.tuner_id = tuner_id

    tuner.search(
        x=x,
        y=y,
        epochs=2,
        steps_per_epoch=20,
        validation_steps=10,
        validation_data=(val_x, val_y),
    )
    return tuner
    def testCloudTunerDatasets(self):
        """Test case to configure Tuner with tf.data.Dataset as input data."""
        train_dataset = tf.data.Dataset.from_tensor_slices(
            (self._x, self._y)).batch(128).cache().prefetch(1000)
        eval_dataset = tf.data.Dataset.from_tensor_slices(
            (self._val_x, self._val_y)).batch(128).cache().prefetch(1000)

        study_id = '{}_dataset'.format(_STUDY_ID_BASE)
        tuner = cloudtuner.CloudTuner(_build_model,
                                      project_id=_PROJECT_ID,
                                      region=_REGION,
                                      objective='acc',
                                      hyperparameters=_HPS,
                                      study_id=study_id,
                                      max_trials=5,
                                      directory=os.path.join(
                                          self.get_temp_dir(), study_id))

        self._assert_output(
            tuner.search_space_summary,
            r'.*Search space summary(?=.*learning_rate \(Float\))'
            r'(?=.*num_layers \(Int\).*)')

        tuner.search(x=train_dataset,
                     epochs=2,
                     steps_per_epoch=20,
                     validation_steps=10,
                     validation_data=eval_dataset)

        self._assert_results_summary(tuner.results_summary)
    def testCloudTunerHyperparameters(self):
        """Test case to configure Tuner with HyperParameters object."""
        study_id = '{}_hyperparameters'.format(_STUDY_ID_BASE)
        tuner = cloudtuner.CloudTuner(_build_model,
                                      project_id=_PROJECT_ID,
                                      region=_REGION,
                                      objective='acc',
                                      hyperparameters=_HPS,
                                      max_trials=5,
                                      study_id=study_id,
                                      directory=os.path.join(
                                          self.get_temp_dir(), study_id))

        # "Search space summary" comes first, but the order of
        # "learning_rate (Float)" and "num_layers (Int)" is not deterministic,
        # hence they are wrapped as look-ahead assertions in the regex.
        self._assert_output(
            tuner.search_space_summary,
            r'.*Search space summary(?=.*learning_rate \(Float\))'
            r'(?=.*num_layers \(Int\).*)')

        tuner.search(x=self._x,
                     y=self._y,
                     epochs=2,
                     steps_per_epoch=20,
                     validation_steps=10,
                     validation_data=(self._val_x, self._val_y))

        self._assert_results_summary(tuner.results_summary)
    def testCloudTunerStudyConfig(self):
        """Test case to configure Tuner with StudyConfig object."""
        # Configure the search space. Specification:
        # https://cloud.google.com/ai-platform/optimizer/docs/reference/rest/v1/projects.locations.studies#StudyConfig
        study_config = {
            'metrics': [{
                'goal': 'MAXIMIZE',
                'metric': 'acc'
            }],
            'parameters': [{
                'discrete_value_spec': {
                    'values': [0.0001, 0.001, 0.01]
                },
                'parameter': 'learning_rate',
                'type': 'DISCRETE'
            }, {
                'integer_value_spec': {
                    'max_value': 10,
                    'min_value': 2
                },
                'parameter': 'num_layers',
                'type': 'INTEGER'
            }, {
                'discrete_value_spec': {
                    'values': [32, 64, 96, 128]
                },
                'parameter': 'units',
                'type': 'DISCRETE'
            }],
            'algorithm':
            'ALGORITHM_UNSPECIFIED',
            'automatedStoppingConfig': {
                'decayCurveStoppingConfig': {
                    'useElapsedTime': True
                }
            }
        }

        study_id = '{}_study_config'.format(_STUDY_ID_BASE)
        tuner = cloudtuner.CloudTuner(_build_model,
                                      project_id=_PROJECT_ID,
                                      region=_REGION,
                                      study_config=study_config,
                                      study_id=study_id,
                                      max_trials=5,
                                      directory=os.path.join(
                                          self.get_temp_dir(), study_id))

        self._assert_output(
            tuner.search_space_summary,
            r'.*Search space summary(?=.*learning_rate \(Choice\))'
            r'(?=.*num_layers \(Int\))(?=.*units \(Choice\))')

        tuner.search(x=self._x,
                     y=self._y,
                     epochs=2,
                     steps_per_epoch=20,
                     validation_steps=10,
                     validation_data=(self._val_x, self._val_y))

        self._assert_results_summary(tuner.results_summary)