internal cleanup

PiperOrigin-RevId: 319343653

Author: ucdmkt

tfx/components/tuner/executor.py

@@ -85,12 +85,11 @@ def search(input_dict: Dict[Text, List[types.Artifact]],
 
   tuner_fn_result = tuner_fn(fn_args)
   result = tuner_fn_result.tuner
-  fit_kwargs = tuner_fn_result.fit_kwargs
 
   # TODO(b/156966497): set logger for printing.
   result.search_space_summary()
   logging.info('Start tuning... Tuner ID: %s', result.tuner_id)
-  result.search(**fit_kwargs)
+  result.search(**tuner_fn_result.fit_kwargs)
   logging.info('Finished tuning... Tuner ID: %s', result.tuner_id)
   result.results_summary()
 

tfx/extensions/google_cloud_ai_platform/tuner/executor.py

@@ -154,7 +154,10 @@ def run_chief_oracle():
       """Invoke chief oracle, and listen to the open port."""
       logging.info('chief_oracle() starting...')
 
-      # Per specification, KerasTuner's behavior is controlled by env variables.
+      # Per KerasTuner's specification, configuration of chief oracle is set
+      # by environment variables. This only affects the current sub-process
+      # which is single-threaded, but not the main process. As such, mutation
+      # of this otherwise global state is safe.
       os.environ['KERASTUNER_ORACLE_IP'] = '0.0.0.0'
       os.environ['KERASTUNER_ORACLE_PORT'] = self._master_port
       os.environ['KERASTUNER_TUNER_ID'] = 'chief'
@@ -198,6 +201,9 @@ def _search(self, input_dict: Dict[Text, List[types.Artifact]],
             input_dict, exec_properties)
 
       # If distributed, both master and worker need to know where the oracle is.
+      # Per KerasTuner's interface, it is configured through env variables.
+      # This only affects the current main process, which is designed to be
+      # single-threaded.
       os.environ['KERASTUNER_ORACLE_IP'] = self._master_addr
       os.environ['KERASTUNER_ORACLE_PORT'] = self._master_port