a Sic7@s~dZddlZddlZddlmmZddiZdZddZ dd Z d d Z d d Z ddZ ddZddZddZddZdS)aUtilities that help manage directory path in distributed settings. In multi-worker training, the need to write a file to distributed file location often requires only one copy done by one worker despite many workers that are involved in training. The option to only perform saving by chief is not feasible for a couple of reasons: 1) Chief and workers may each contain a client that runs the same piece of code and it's preferred not to make any distinction between the code run by chief and other workers, and 2) saving of model or model's related information may require SyncOnRead variables to be read, which needs the cooperation of all workers to perform all-reduce. This set of utility is used so that only one copy is written to the needed directory, by supplying a temporary write directory path for workers that don't need to save, and removing the temporary directory once file writing is done. Example usage: ``` # Before using a directory to write file to. self.log_write_dir = write_dirpath(self.log_dir, get_distribution_strategy()) # Now `self.log_write_dir` can be safely used to write file to. ... # After the file is written to the directory. remove_temp_dirpath(self.log_dir, get_distribution_strategy()) ``` Experimental. API is subject to change. NzMetadata-FlavorGoogleGCE_METADATA_IPcCs|jj}dt|S)N workertemp_)extended_task_idstr)strategytask_idr c/var/www/html/django/DPS/env/lib/python3.9/site-packages/keras/distribute/distributed_file_utils.py_get_base_dirpath8sr cCs|t|SN)endswithr dirpathrr r r _is_temp_dir=srcCs4t||r|}ntj|t|}tjj||Sr ) rospathjoinr tfiogfilemakedirs)rrtemp_dirr r r _get_temp_dirAs  rcCsB|durtj}|dur|S|js,|S|jjr8|St||S)aoReturns the writing dir that should be used to save file distributedly. `dirpath` would be created if it doesn't exist. Args: dirpath: Original dirpath that would be used without distribution. strategy: The tf.distribute strategy object currently used. Returns: The writing dir path that should be used to save with distribution. N)r distribute get_strategyr_in_multi_worker_modeshould_checkpointrrr r r write_dirpathJs   rcCsJ|durtj}|durdS|jrF|jjsFtjjj t ||dS)zRemoves the temp path after writing is finished. Args: dirpath: Original dirpath that would be used without distribution. strategy: The tf.distribute strategy object currently used. N) rrrrrrcompatv1rDeleteRecursivelyrrr r r remove_temp_dirpathfs r#cCs,tj|}tj|}tjt|||S)aReturns the writing file path to be used to save file distributedly. Directory to contain `filepath` would be created if it doesn't exist. Args: filepath: Original filepath that would be used without distribution. strategy: The tf.distribute strategy object currently used. Returns: The writing filepath that should be used to save file with distribution. )rrdirnamebasenamerr)filepathrrbaser r r write_filepaths  r(cCsttj||dS)zRemoves the temp path for file after writing is finished. Args: filepath: Original filepath that would be used without distribution. strategy: The tf.distribute strategy object currently used. N)r#rrr$)r&rr r r remove_temp_dir_with_filepathsr)cCsPdtjtd}z tjd|dftdd}|jWStjjyJYdS0dS) z9Detect whether the current running environment is on GCP.zhttp://zmetadata.google.internalz%s/computeMetadata/v1/%szinstance/hostname)headerstimeoutFN) renvironget_GCE_METADATA_URL_ENV_VARIABLErequestsGCP_METADATA_HEADER status_code exceptionsRequestException)gce_metadata_endpointresponser r r _on_gcpsr7cCstrtjdsdSdS)NTPUTF)r7rconfiglist_physical_devicesr r r r %support_on_demand_checkpoint_callbacksr;)__doc__rr0tensorflow.compat.v2r v2rr1r/r rrrr#r(r)r7r;r r r r s