1# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2# SPDX-License-Identifier: Apache-2.0
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8# http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15
16"""
17Hydra ConfigSource plugin for Multi-Storage Client.
18
19This module provides a ConfigSource implementation that allows Hydra to load
20configuration files from remote storage systems using Multi-Storage Client.
21"""
22
23import logging
24
25from hydra.core.config_search_path import ConfigSearchPath, SearchPathElement
26from hydra.core.object_type import ObjectType
27from hydra.plugins.config_source import ConfigLoadError, ConfigResult, ConfigSource
28from hydra.plugins.search_path_plugin import SearchPathPlugin
29from omegaconf import OmegaConf
30
31import multistorageclient as msc
32from multistorageclient.shortcuts import resolve_storage_client
33from multistorageclient.types import MSC_PROTOCOL, MSC_PROTOCOL_NAME
34from multistorageclient.utils import join_paths
35
36logger = logging.getLogger(__name__)
37
38
[docs]
39class MSCConfigSource(ConfigSource):
40 """
41 A Hydra :py:class:`hydra.plugins.config_source.ConfigSource` that uses Multi-Storage Client to read configuration files from remote storage systems.
42
43 Supports loading configs from S3, GCS, Azure Blob Storage, and other MSC-supported storage backends.
44 Must be used in conjunction with :py:class:`MSCSearchPathPlugin`.
45 """
46
47 def __init__(self, provider: str, path: str) -> None:
48 """
49 Initialize the MSC ConfigSource.
50
51 :param provider: The provider name (should be ``main`` if the user specifies config path, or ``msc-universal`` if the user doesn't specify a config path).
52 :param path: The base path for this source. It should be a full MSC URL (e.g., ``msc://dev/configs``) if the user specifies config path, or ``msc://`` if the user doesn't specify a config path and the universal MSC source is used.
53 """
54 if path.find("://") == -1:
55 path = f"{MSC_PROTOCOL}{path}"
56 super().__init__(provider=provider, path=path)
57
58 # Store the base URL for resolving relative paths
59 self.base_url = path
60 logger.debug(f"Initialized MSCConfigSource with base URL: {self.base_url}")
61
[docs]
62 @staticmethod
63 def scheme() -> str:
64 """
65 Return the URL scheme for this ConfigSource.
66
67 :return: The scheme string 'msc'.
68 """
69 return MSC_PROTOCOL_NAME
70
71 def _resolve_full_url(self, config_path: str) -> str:
72 """
73 Convert config_path to a full ``msc://`` URL.
74
75 :param config_path: Either a relative path, config group reference (e.g., ``database: postgres``), or full ``msc://`` URL.
76 :return: Full ``msc://`` URL that can be passed to ``multistorageclient.open()``.
77 """
78 if config_path.startswith(MSC_PROTOCOL):
79 return config_path
80
81 # Handle Hydra defaults syntax: "group: config" -> "group/config"
82 if ": " in config_path:
83 group, config_name = config_path.split(": ", 1)
84 config_path = f"{group}/{config_name}"
85
86 # Relative path - join with base URL using MSC's utility
87 return join_paths(self.base_url, config_path)
88
[docs]
89 def load_config(self, config_path: str) -> ConfigResult:
90 """
91 Load a configuration file from MSC storage.
92
93 :param config_path: Relative path to the config file, or full ``msc://`` URL.
94 :return: The loaded configuration.
95 :raises ConfigLoadError: If the config file cannot be loaded.
96 """
97 full_url = self._resolve_full_url(config_path)
98 full_url = self._normalize_file_name(full_url)
99
100 try:
101 with msc.open(full_url, "r") as f:
102 header_text = f.read(512)
103 header = ConfigSource._get_header_dict(header_text)
104 f.seek(0)
105 cfg = OmegaConf.load(f)
106 return ConfigResult(
107 provider=self.provider,
108 path=f"{self.scheme()}://{self.path}",
109 config=cfg,
110 header=header,
111 )
112 except Exception as e:
113 raise ConfigLoadError(f"Failed to load config from {full_url}: {e}")
114
[docs]
115 def available(self) -> bool:
116 """
117 Check if the MSC config source is available.
118
119 :return: ``True`` if the MSC config source can be accessed, ``False`` otherwise.
120 """
121 try:
122 # Try to resolve the base URL to see if MSC can handle it
123 resolve_storage_client(self.base_url)
124 return True
125 except Exception:
126 logger.exception("MSC config source not available")
127 return False
128
[docs]
129 def is_group(self, config_path: str) -> bool:
130 """
131 Check if the given path is a group (directory).
132
133 :param config_path: Relative path or full ``msc://`` URL to check.
134 :return: ``True`` if the path is a directory, ``False`` otherwise.
135 """
136 full_url = self._resolve_full_url(config_path)
137
138 # Ensure path ends with "/" for directory check
139 full_url = full_url.rstrip("/") + "/"
140
141 try:
142 # Use msc.info() directly to check if path is a directory
143 metadata = msc.info(full_url)
144 return metadata.type == "directory"
145 except FileNotFoundError:
146 return False
147 except Exception:
148 return False
149
[docs]
150 def is_config(self, config_path: str) -> bool:
151 """
152 Check if the given path is a config file.
153
154 :param config_path: Relative path or full ``msc://`` URL to check.
155 :return: ``True`` if the path is a config file, ``False`` otherwise.
156 """
157 # If there's a directory with the same name, directory takes precedence
158 if self.is_group(config_path):
159 return False
160
161 full_url = self._resolve_full_url(config_path)
162 full_url = self._normalize_file_name(full_url)
163
164 try:
165 return msc.is_file(full_url)
166 except Exception:
167 return False
168
[docs]
169 def list(self, config_path: str, results_filter: ObjectType | None) -> list[str]:
170 """
171 List items under the specified config path.
172
173 :param config_path: Relative path or full ``msc://`` URL to list.
174 :param results_filter: Optional filter for the results.
175 :return: List of config names and group names under the specified path.
176 """
177 full_url = self._resolve_full_url(config_path)
178 files: list[str] = []
179
180 try:
181 # Use MSC to resolve client and list items directly
182 # In this case, client.list() is simpler than msc.list() because
183 # msc.list() returns keys with full msc:// URLs which would require
184 # more complex path manipulation
185 client, path = resolve_storage_client(full_url)
186
187 # Add trailing slash to ensure we're listing directory contents
188 list_path = path.rstrip("/") + "/" if path else ""
189
190 # Get all items under this path
191 for item in client.list(path=list_path, include_directories=True):
192 # Get the relative path from the base path
193 if item.key.startswith(list_path):
194 relative_path = item.key[len(list_path) :]
195 elif path and item.key.startswith(path + "/"):
196 relative_path = item.key[len(path + "/") :]
197 else:
198 continue
199
200 # Skip empty paths or the directory itself
201 if not relative_path or relative_path.endswith("/"):
202 continue
203
204 # Get just the immediate file/directory name (no nested paths)
205 file_name = relative_path.split("/")[0]
206 if not file_name:
207 continue
208
209 # Build the full config path for this item
210 item_path = join_paths(config_path, file_name) if config_path else file_name
211
212 self._list_add_result(
213 files=files,
214 file_path=item_path,
215 file_name=file_name,
216 results_filter=results_filter,
217 )
218
219 except Exception:
220 logger.exception(f"Failed to list MSC path '{full_url}'")
221 # Return empty list if we can't list the directory
222
223 return sorted(set(files))
224
225 def __repr__(self) -> str:
226 return f"MSCConfigSource(provider={self.provider}, path={self.scheme()}://{self.path})"
227
228
[docs]
229class MSCSearchPathPlugin(SearchPathPlugin):
230 """
231 A Hydra :py:class:`hydra.plugins.search_path_plugin.SearchPathPlugin` that enables MSC support.
232
233 Fixes MSC URL mangling issues and ensures MSC sources are available for config loading.
234 """
235
[docs]
236 def manipulate_search_path(self, search_path: ConfigSearchPath) -> None:
237 """
238 Enable MSC support by fixing mangled URLs and adding universal MSC source.
239
240 Performs two operations:
241
242 1. **Fixes mangled MSC URLs**: CLI path normalization can mangle ``msc://dev/configs`` to ``/current/dir/msc:/dev/configs``.
243 2. **Adds universal MSC source**: Ensures :py:class:`MSCConfigSource` is available to handle any ``msc://`` URLs in config defaults.
244
245 :param search_path: The :py:class:`hydra.core.config_search_path.ConfigSearchPath` to manipulate.
246 """
247 path_elements = search_path.get_path()
248
249 # Step 1: Fix any mangled MSC URLs from CLI
250 for i, element in enumerate(path_elements):
251 path = element.path
252
253 # Detect mangled MSC URLs: contains "msc:/" but doesn't start with "msc://"
254 if path and "msc:/" in path and not path.startswith(MSC_PROTOCOL):
255 # Extract the MSC URL from the mangled path
256 msc_start = path.find("msc:/")
257 msc_part = path[msc_start:] # Everything from "msc:/" onwards
258
259 # Fix the missing slash: "msc:/profile" -> "msc://profile"
260 if msc_part.startswith("msc:/") and not msc_part.startswith(MSC_PROTOCOL):
261 fixed_url = msc_part.replace("msc:/", MSC_PROTOCOL, 1)
262 else:
263 fixed_url = msc_part
264
265 # Replace the element with a new one containing the fixed URL
266 path_elements[i] = SearchPathElement(element.provider, fixed_url)
267
268 logger.debug(f"Fixed mangled MSC URL: {path} → {fixed_url}")
269
270 # Step 2: Ensure there's a universal MSC source for handling any msc:// URLs
271 # Check if there's already an msc:// entry in the search path
272 has_msc_source = any(element.path and element.path.startswith(MSC_PROTOCOL) for element in path_elements)
273
274 if not has_msc_source:
275 # Add a universal MSC source that can handle any msc:// URL
276 # Use a generic base that MSCConfigSource can resolve dynamically
277 search_path.append("msc-universal", MSC_PROTOCOL)
278 logger.debug("Added universal MSC source to search path")