custom ConfigSource plugin appears to load configuration twice; any way to disable being used in 'get_mode' call?

[LrWm3]

I am working on writing two ConfigSource plugins; one for AWS Secrets Manager and Weights and Biases. I was able to get these plugins working, but I noticed that they are run twice by Hydra.

These plugins are quite slow compared to the file system plugins Hydra ships with, since they make network calls. Is there any way to only run these plugins during the main hydra configuration loading and not run them the first time?

Checking the source, it appears the first time it is called is during 'get_mode'

hydra/hydra/_internal/hydra.py

Lines 81 to 96 in 2f76287

	def get_mode(
	self,
	config_name: Optional[str],
	overrides: List[str],
	) -> Any:
	try:
	cfg = self.compose_config(
	config_name=config_name,
	overrides=overrides,
	with_log_configuration=False,
	run_mode=RunMode.MULTIRUN,
	validate_sweep_overrides=False,
	)
	return cfg.hydra.mode
	except Exception:
	return None

However, this is only to get the cfg.hydra.mode value.

Is there is a way to specify ConfigSource plugins do not support this value to avoid running it twice? This would significantly speed up the time it takes to start the application.

edit: I guess its not a significant speed-up if local caching is used but its awkward to explain why it is happening to devs and work around from inside the plugin.

[odelalleau]

Not sure, but maybe you can use a cache in your plugins so that the second call is much faster?

[LrWm3]

I am doing some IoT work, so I have implemented filesystem caching which can work if internet becomes unavailable. Additionally, I index remote resources to reduce the API calls to the absolute minimum and it does speed things up quite a bit. From 33 seconds to about 4 seconds all in (both calls)

Then I force reuse of the filesystem cache on the second run with no API calls. if I kept everything in memory it would be faster but I have not implemented that as I wanted to see if there was another way first. So it ends up being another 500ms for the second call.

2023-12-06 12:51:17.939 | DEBUG    | hydra_plugins.aws_secrets_config_source:__init__:57 - Initializing AWS Secrets Config Source
2023-12-06 12:51:17.940 | DEBUG    | hydra_plugins.wandb_config_source:__init__:92 - WandB artifact loading
2023-12-06 12:51:20.500 | DEBUG    | hydra_plugins.wandb_config_source:__del__:394 - WandB artifact cache clean-up disabled
2023-12-06 12:51:20.500 | DEBUG    | hydra_plugins.aws_secrets_config_source:__del__:185 - Cleaning up AWS Secrets Config Source
2023-12-06 12:51:21.392 | DEBUG    | hydra_plugins.aws_secrets_config_source:__init__:57 - Initializing AWS Secrets Config Source
2023-12-06 12:51:21.393 | DEBUG    | hydra_plugins.wandb_config_source:__init__:92 - WandB artifact loading
2023-12-06 12:51:21.897 | DEBUG    | hydra_plugins.wandb_config_source:__del__:394 - WandB artifact cache clean-up disabled
2023-12-06 12:51:21.897 | DEBUG    | hydra_plugins.aws_secrets_config_source:__del__:185 - Cleaning up AWS Secrets Config Source

However it is a bit weird because hydra even has a 'caching' config repository. but I don't get to take advantage of it because hydra seems to completely reload the config from scratch after this get_mode call. Maybe there is a technical reason for this, I am still not very familiar with how hydra works.

hydra/hydra/_internal/config_repository.py

Lines 328 to 366 in 2f76287

	class CachingConfigRepository(IConfigRepository):
	def __init__(self, delegate: IConfigRepository):
	# copy the underlying repository to avoid mutating it with initialize_sources()
	self.delegate = copy.deepcopy(delegate)
	self.cache: Dict[str, Optional[ConfigResult]] = {}
	def get_schema_source(self) -> ConfigSource:
	return self.delegate.get_schema_source()
	def initialize_sources(self, config_search_path: ConfigSearchPath) -> None:
	self.delegate.initialize_sources(config_search_path)
	# not clearing the cache.
	# For the use case this is used, the only thing in the cache is the primary config
	# and we want to keep it even though we re-initialized the sources.
	def load_config(self, config_path: str) -> Optional[ConfigResult]:
	cache_key = f"config_path={config_path}"
	if cache_key in self.cache:
	return self.cache[cache_key]
	else:
	ret = self.delegate.load_config(config_path=config_path)
	self.cache[cache_key] = ret
	return ret
	def group_exists(self, config_path: str) -> bool:
	return self.delegate.group_exists(config_path=config_path)
	def config_exists(self, config_path: str) -> bool:
	return self.delegate.config_exists(config_path=config_path)
	def get_group_options(
	self, group_name: str, results_filter: Optional[ObjectType] = ObjectType.CONFIG
	) -> List[str]:
	return self.delegate.get_group_options(
	group_name=group_name, results_filter=results_filter
	)
	def get_sources(self) -> List[ConfigSource]:
	return self.delegate.get_sources()

Just wanted to see if anyone knew of other options before going further with caching based approaches.

[LrWm3]

Or maybe I should be wrapping my plugin with a global instance of this caching config repository.... I didn't think of this until just now... 😢

edit: actually that doesn't quite make sense since I am doing a ConfigPlugin not a Repository 😅

[odelalleau]

I have no idea what this CachingConfigRepository does so I'm afraid I can't provide advice on how to use it, sorry.

[LrWm3]

Its okay. I added in memory caching to my plugins and have gotten the second execution down to 200ms. I think this will be good enough.

Wanted to have the cache to be static to survive plugin deletion so I used a global decorator. I've included the code snippet below in case someone else has the same issue.

could be improved as a class-level annotation. i also am not sure if it needs to be global or if defining it in aws_secrets_manager_global_cache would work.

def aws_secrets_manager_global_cache(func):
    """
    Hydra calls the ConfigSource plugin multiple times, and 
    creates and deletes the plugins several times as well 
    
    So we need to cache everything to avoid taking a long time
    to load the configuration.
    """
    _g_aws_secrets_manager_cache = {}
    def wrapper(*args, **kwargs):
        cache_key = func.__name__ + str(args) + str(kwargs)
        if cache_key not in _g_aws_secrets_manager_cache:
            _g_aws_secrets_manager_cache[cache_key] = func(*args, **kwargs)
        return _g_aws_secrets_manager_cache[cache_key]
    return wrapper

# then decorate each function to cache

class AWSSecretsConfigSource(ConfigSource):

    @aws_secrets_manager_global_cache
    def load_config(self, config_path: str) -> ConfigResult:

    # ... rest ...

far from my best work but does seem to do the trick.