refactor fetch_and_render into components

Author: simon-mo

Diff for: .gitignore

@@ -2,6 +2,7 @@
 bazel_events
 travis_events
 results.db
+cache_dir/
 
 .vscode/
 

Diff for: ray_ci_tracker/__init__.py

@@ -0,0 +1,121 @@
+import asyncio
+import functools
+import os
+from datetime import datetime
+from functools import wraps
+from itertools import chain
+from pathlib import Path
+from subprocess import PIPE
+from typing import List, Optional, Tuple
+
+import aiofiles
+import click
+import httpx
+import ujson as json
+from dotenv import load_dotenv
+from tqdm.asyncio import tqdm_asyncio
+
+from ray_ci_tracker.interfaces import (
+    BuildkiteArtifact,
+    BuildkiteStatus,
+    BuildResult,
+    GHAJobStat,
+    GHCommit,
+    TestResult,
+    _parse_duration,
+)
+
+
+def retry(func):
+    @wraps(func)
+    async def wrapper(*args, **kwargs):
+        exception = None
+        for _ in range(3):
+            try:
+                return await func(*args, **kwargs)
+            except Exception as e:
+                exception = e
+        else:
+            raise exception
+
+    return wrapper
+
+
+def run_as_sync(async_func):
+    @wraps(async_func)
+    def wrapper(*args, **kwargs):
+        asyncio.run(async_func(*args, **kwargs))
+
+    return wrapper
+
+
+async def get_or_fetch(
+    cache_path: Path, *, use_cached: bool, result_cls, many: bool, async_func
+):
+    if not use_cached or not cache_path.exists():
+        result = await async_func()
+        if result is None:
+            return
+        cache_path.parent.mkdir(parents=True, exist_ok=True)
+        async with aiofiles.open(cache_path, "w") as f:
+            if result_cls and many:
+                content = json.dumps([r.to_dict() for r in result])
+            elif result_cls and not many:
+                content = json.dumps(result.to_dict())
+            else:
+                content = json.dumps(result)
+            await f.write(content)
+        return result
+    else:
+        async with aiofiles.open(cache_path) as f:
+            content = await f.read()
+            if result_cls and many:
+                loaded = json.loads(content)
+                return [result_cls.from_dict(r) for r in loaded]
+            elif result_cls and not many:
+                return result_cls.from_json(content)
+            else:
+                return json.loads(content)
+
+
+def _yield_test_result(bazel_log_path):
+    # Gather the known flaky set
+    flaky_tests = set()
+    with open(bazel_log_path) as f:
+        for line in f:
+            loaded = json.loads(line)
+            if "targetConfigured" in loaded["id"] and "tag" in loaded["configured"]:
+                test_name = loaded["id"]["targetConfigured"]["label"]
+                if "flaky" in loaded["configured"]["tag"]:
+                    flaky_tests.add(test_name)
+
+    with open(bazel_log_path) as f:
+        for line in f:
+            loaded = json.loads(line)
+            if "testSummary" in loaded:
+                test_summary = loaded["testSummary"]
+
+                name = loaded["id"]["testSummary"]["label"]
+                status = test_summary["overallStatus"]
+                if status in {"FAILED", "TIMEOUT", "NO_STATUS"}:
+                    status = "FAILED"
+                duration_s = float(test_summary["totalRunDurationMillis"]) / 1e3
+                yield TestResult(name, status, duration_s, name in flaky_tests)
+
+
+def _process_single_build(dir_name) -> BuildResult:
+    with open(dir_name / "metadata.json") as f:
+        metadata = json.load(f)
+
+    return BuildResult(
+        sha=metadata["build_env"]["TRAVIS_COMMIT"],
+        job_url=metadata["build_env"]["TRAVIS_JOB_WEB_URL"],
+        os=metadata["build_env"]["TRAVIS_OS_NAME"],
+        build_env=metadata["build_config"]["config"]["env"],
+        job_id=os.path.split(dir_name)[-1],
+        results=list(
+            chain.from_iterable(
+                _yield_test_result(log) for log in dir_name.glob("bazel_log.*")
+            )
+        ),
+    )

Diff for: ray_ci_tracker/common.py

@@ -0,0 +1,196 @@
+import asyncio
+import functools
+import os
+from itertools import chain
+from pathlib import Path
+from typing import List
+
+import aiofiles
+import httpx
+from tqdm.asyncio import tqdm_asyncio
+
+from ray_ci_tracker.common import _process_single_build, get_or_fetch, retry
+from ray_ci_tracker.interfaces import BuildkiteArtifact, BuildkiteStatus, BuildResult
+
+GRAPHQL_QUERY = """
+query AllPipelinesQuery {
+  pipeline(slug: "ray-project/ray-builders-branch") {
+    builds(branch: "master", commit: "COMMIT_PLACEHODLER") {
+      count
+      edges {
+        node {
+          createdAt
+          startedAt
+          finishedAt
+          number
+          jobs(first: 100) {
+            edges {
+              node {
+                ... on JobTypeCommand {
+                  uuid
+                  label
+                  passed
+                  state
+                  url
+                  build {
+                    commit
+                  }
+                  createdAt
+                  runnableAt
+                  startedAt
+                  finishedAt
+                  artifacts(first: 100) {
+                    edges {
+                      node {
+                        downloadURL
+                        path
+                      }
+                    }
+                  }
+                }
+              }
+            }
+          }
+        }
+      }
+    }
+  }
+}
+"""
+
+
+class BuildkiteSource:
+    @staticmethod
+    async def fetch_all(cache_path: Path, cached_buildkite, commits):
+        print("Downloading Buildkite Status (Jobs)")
+        concurrency_limiter = asyncio.Semaphore(20)
+        buildkite_jsons = await tqdm_asyncio.gather(
+            *[
+                get_or_fetch(
+                    cache_path / f"bk_jobs/{commit.sha}/http_resp.json",
+                    use_cached=cached_buildkite,
+                    result_cls=None,
+                    many=False,
+                    async_func=functools.partial(
+                        BuildkiteSource.get_buildkite_job_status,
+                        commit_sha=commit.sha,
+                        concurrency_limiter=concurrency_limiter,
+                    ),
+                )
+                for commit in commits
+            ]
+        )
+        buildkite_parsed: List[BuildkiteStatus] = await asyncio.gather(
+            *[
+                get_or_fetch(
+                    cache_path / f"bk_jobs/{commit.sha}/parsed.json",
+                    use_cached=cached_buildkite,
+                    result_cls=BuildkiteStatus,
+                    many=True,
+                    async_func=functools.partial(
+                        BuildkiteSource.parse_buildkite_build_json,
+                        resp_json,
+                    ),
+                )
+                for commit, resp_json in zip(commits, buildkite_jsons)
+            ]
+        )
+        print("Downloading Buildkite Status (macOS Bazel Events)")
+        macos_bazel_events = await tqdm_asyncio.gather(
+            *[
+                get_or_fetch(
+                    cache_path
+                    / f"bazel_cached/{status.commit}/mac_result_{status.job_id}.json",
+                    use_cached=cached_buildkite,
+                    result_cls=BuildResult,
+                    many=False,
+                    async_func=functools.partial(
+                        BuildkiteSource.get_buildkite_artifact,
+                        dir_prefix=cache_path,
+                        artifacts=status.artifacts,
+                        concurrency_limiter=asyncio.Semaphore(50),
+                    ),
+                )
+                for status in chain.from_iterable(buildkite_parsed)
+                if len(status.artifacts) > 0
+            ]
+        )
+        return list(chain.from_iterable(buildkite_parsed)), macos_bazel_events
+
+    @staticmethod
+    @retry
+    async def get_buildkite_job_status(
+        commit_sha, concurrency_limiter: asyncio.Semaphore
+    ):
+        http_client = httpx.AsyncClient(timeout=httpx.Timeout(60))
+        async with concurrency_limiter, http_client:
+            resp = await http_client.post(
+                "https://graphql.buildkite.com/v1",
+                headers={"Authorization": f"Bearer {os.environ['BUILDKITE_TOKEN']}"},
+                json={"query": GRAPHQL_QUERY.replace("COMMIT_PLACEHODLER", commit_sha)},
+            )
+        resp.raise_for_status()
+        return resp.json()
+
+    @staticmethod
+    async def parse_buildkite_build_json(
+        resp_json: dict,
+    ) -> List[BuildkiteStatus]:
+        builds = resp_json["data"]["pipeline"]["builds"]["edges"]
+
+        statuses = []
+        for build in builds:
+            jobs = build["node"]["jobs"]["edges"]
+            for job in jobs:
+                actual_job = job["node"]
+                job_id = actual_job["uuid"]
+                sha = actual_job["build"]["commit"]
+
+                artifacts = []
+                for artifact in actual_job["artifacts"]["edges"]:
+                    url = artifact["node"]["downloadURL"]
+                    path = artifact["node"]["path"]
+                    filename = os.path.split(path)[1]
+                    on_disk_path = f"bazel_events/master/{sha}/{job_id}/{filename}"
+                    artifacts.append(
+                        BuildkiteArtifact(url=url, bazel_events_path=on_disk_path)
+                    )
+
+                status = BuildkiteStatus(
+                    job_id=job_id,
+                    label=actual_job["label"],
+                    passed=actual_job["passed"],
+                    state=actual_job["state"],
+                    url=actual_job["url"],
+                    commit=sha,
+                    startedAt=actual_job["startedAt"],
+                    finished_at=actual_job["finishedAt"],
+                    artifacts=artifacts,
+                )
+                statuses.append(status)
+        return statuses
+
+    @staticmethod
+    @retry
+    async def get_buildkite_artifact(
+        dir_prefix: Path,
+        artifacts: List[BuildkiteArtifact],
+        concurrency_limiter: asyncio.Semaphore,
+    ) -> BuildResult:
+        assert len(artifacts)
+
+        bazel_events_dir = None
+        async with concurrency_limiter, httpx.AsyncClient() as client:
+            for artifact in artifacts:
+                path = dir_prefix / artifact.bazel_events_path
+                path.parent.mkdir(exist_ok=True, parents=True)
+                bazel_events_dir = path.parent
+                async with client.stream(
+                    "GET", artifact.url
+                ) as response, aiofiles.open(path, "wb") as f:
+                    response.raise_for_status()
+                    async for chunk in response.aiter_bytes():
+                        await f.write(chunk)
+
+        assert bazel_events_dir is not None
+        return _process_single_build(bazel_events_dir)