[feat] Add cloud parameter

collect_single/__main__.py

@@ -51,6 +51,7 @@ async def collect_and_sync(send_on_failed: bool) -> bool:
             async with CollectAndSync(
                 redis=redis,
                 tenant_id=args.tenant_id,
+                cloud=args.cloud,
                 account_id=args.account_id,
                 job_id=args.job_id,
                 core_args=core_args,
@@ -97,6 +98,7 @@ def main() -> None:
     parser.add_argument("--job-id", required=True, help="Job Id of the coordinator")
     parser.add_argument("--tenant-id", required=True, help="Id of the tenant")
     parser.add_argument("--account-id", help="Id of the account")
+    parser.add_argument("--cloud", help="Cloud provider.")
     parser.add_argument("--redis-url", default="redis://localhost:6379/0", help="Redis host.")
     parser.add_argument("--redis-password", default=os.environ.get("REDIS_PASSWORD"), help="Redis password")
     parser.add_argument("--push-gateway-url", help="Prometheus push gateway url")

collect_single/collect_and_sync.py

@@ -41,6 +41,7 @@ def __init__(
         *,
         redis: Redis,  # type: ignore
         tenant_id: str,
+        cloud: Optional[str],
         account_id: Optional[str],
         job_id: str,
         core_args: List[str],
@@ -51,6 +52,7 @@ def __init__(
     ) -> None:
         self.redis = redis
         self.tenant_id = tenant_id
+        self.cloud = cloud
         self.account_id = account_id
         self.job_id = job_id
         self.core_args = ["fixcore", "--no-scheduling", "--ignore-interrupted-tasks"] + core_args
@@ -127,7 +129,7 @@ async def post_process(self) -> Tuple[Json, List[str]]:
         # post process the data, if something has been collected
         if account_info and (account_id := self.account_id):
             # synchronize the security section
-            benchmarks = await self.core_client.list_benchmarks()
+            benchmarks = await self.core_client.list_benchmarks(providers=[self.cloud] if self.cloud else None)
             if benchmarks:
                 await self.core_client.create_benchmark_reports(account_id, benchmarks, self.task_id)
             # create metrics

collect_single/core_client.py

@@ -3,7 +3,7 @@
 import asyncio
 import logging
 import uuid
-from typing import Any, Optional, Dict, List, Set
+from typing import Any, Optional, Dict, List, Set, cast
 
 from fixcloudutils.types import Json, JsonElement
 from fixclient import Subscriber
@@ -57,8 +57,15 @@ async def workflow_log(self, task_id: str, limit: int = 100) -> List[str]:
             if info != "No error messages for this run."
         ]
 
-    async def list_benchmarks(self) -> List[str]:
-        return [cfg async for cfg in self.client.cli_execute("report benchmark list")]
+    async def list_benchmarks(self, *, providers: Optional[List[str]] = None) -> List[str]:
+        params = dict(ids_only="true")
+        if providers:
+            params["providers"] = " ".join(providers)
+        response = await self.client._get("/report/benchmarks", params)
+        if response.status_code == 200:
+            return cast(List[str], await response.json())
+        else:
+            raise AttributeError(await response.text())
 
     async def create_benchmark_reports(self, account_id: str, benchmarks: List[str], task_id: Optional[str]) -> None:
         bn = " ".join(benchmarks)

tests/conftest.py

@@ -50,6 +50,7 @@ async def collect_and_sync(redis: Redis) -> CollectAndSync:  # type: ignore
     return CollectAndSync(
         redis=redis,
         tenant_id="tenant_id",
+        cloud="aws",
         account_id="account_id",
         job_id="job_id",
         core_args=[],