Attempt to batch config loading for tron deployments

Right now we make at most 2N calls to the Tron API during config
deployments: N to get the current configs and at most N if all services
have changes.

To start, I'd like to reduce this to N by allowing GET /api/config to
return all the configs so that the only requests needed are POSTs for
changed configs. Since I'm a little worried about how tron-pnw-prod will
fare with returning all the configs in a single request, using this new
endpoint is behind a feature toggle. Hopefully this works out and we
don't need to add pagination to the endpoint :)

Depending on how this goes, we can look into batching up the POSTs so
that we can also do that in a single request (or at least <N requests if
we want to chunk things a bit more).

paasta_tools/setup_tron_namespace.py

@@ -23,6 +23,7 @@
 import argparse
 import logging
 import sys
+from typing import Dict
 from typing import List
 
 import ruamel.yaml as yaml
@@ -62,6 +63,13 @@ def parse_args():
     )
     parser.add_argument("-v", "--verbose", action="store_true", default=False)
     parser.add_argument("--dry-run", action="store_true", default=False)
+    parser.add_argument(
+        "--bulk-config-fetch",
+        dest="bulk_config_fetch",
+        action="store_true",
+        default=False,
+        help="Attempt to fetch all configs in bulk rather than one by one",
+    )
     parser.add_argument(
         "--cluster",
         help="Cluster to read configs for. Defaults to the configuration in /etc/paasta",
@@ -162,6 +170,7 @@ def main():
     k8s_enabled_for_cluster = (
         yaml.safe_load(master_config).get("k8s_options", {}).get("enabled", False)
     )
+    new_configs: Dict[str, str] = {}  # service -> new_config
     for service in sorted(services):
         try:
             new_config = tron_tools.create_complete_config(
@@ -171,6 +180,7 @@ def main():
                 k8s_enabled=k8s_enabled_for_cluster,
                 dry_run=args.dry_run,
             )
+            new_configs[service] = new_config
             if args.dry_run:
                 log.info(f"Would update {service} to:")
                 log.info(f"{new_config}")
@@ -187,16 +197,31 @@ def main():
                     for_validation=False,
                 )
                 ensure_service_accounts(job_configs)
+                if not args.bulk_config_fetch:
+                    if client.update_namespace(service, new_config):
+                        updated.append(service)
+                        log.debug(f"Updated {service}")
+                    else:
+                        skipped.append(service)
+                        log.debug(f"Skipped {service}")
 
-                if client.update_namespace(service, new_config):
-                    updated.append(service)
-                    log.debug(f"Updated {service}")
-                else:
-                    skipped.append(service)
-                    log.debug(f"Skipped {service}")
         except Exception:
-            log.exception(f"Update for {service} failed:")
+            log.exception(
+                f"Failed to create service account for {service} (will skip reconfiguring):"
+            )
             failed.append(service)
+            new_configs.pop(service, None)
+
+    if args.bulk_config_fetch:
+        updated_namespaces = client.update_namespaces(new_configs)
+
+        if updated_namespaces:
+            updated = list(updated_namespaces.keys())
+            log.debug(f"Updated {updated}")
+
+        if updated_namespaces != new_configs.keys():
+            skipped = set(new_configs.keys()) - set(updated_namespaces.keys())
+            log.debug(f"Skipped {skipped}")
 
     skipped_report = skipped if args.verbose else len(skipped)
     log.info(

paasta_tools/tron/client.py

@@ -11,6 +11,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import logging
+from typing import Dict
 from urllib.parse import urljoin
 
 import requests
@@ -96,6 +97,38 @@ def update_namespace(self, namespace, new_config, skip_if_unchanged=True):
             },
         )
 
+    def update_namespaces(
+        self, new_configs: Dict[str, str], skip_if_unchanged: bool = True
+    ):
+        """Updates the configuration for a namespace.
+
+        :param namespace: str
+        :param new_config: str, should be valid YAML.
+        :param skip_if_unchanged: boolean. If False, will send the update
+            even if the current config matches the new config.
+        """
+        current_configs: Dict[str, Dict[str, str]] = self._get("/api/config")  # type: ignore  # we don't have a good way to share types between tron/paasta
+        responses: Dict[str, str] = {}
+        for namespace, new_config in new_configs.items():
+            current_config = current_configs.get(namespace, {})
+            if skip_if_unchanged:
+                if yaml.safe_load(new_config) == yaml.safe_load(
+                    current_config["config"]
+                ):
+                    log.debug("No change in config, skipping update.")
+                    continue
+
+            responses[namespace] = self._post(
+                "/api/config",
+                data={
+                    "name": namespace,
+                    "config": new_config,
+                    "hash": current_config["hash"],
+                    "check": 0,
+                },
+            )
+        return responses
+
     def list_namespaces(self):
         """Gets the namespaces that are currently configured."""
         response = self._get("/api")