From 6afa67811c48c08ee26cd7f917c2cd6b786a6e0e Mon Sep 17 00:00:00 2001
From: Cesar Cardoso <cesarcardoso@meta.com>
Date: Thu, 8 Aug 2024 12:29:52 -0700
Subject: [PATCH] Healthcheck analysis class (#2646)

Summary:
Pull Request resolved: https://github.com/facebook/Ax/pull/2646

Add abstract class for health checks

Reviewed By: mpolson64

Differential Revision: D60417357

fbshipit-source-id: 47beb0aeaafd60acd5ba2bf3a8b8cb9674587529
---
 ax/analysis/healthcheck/__init__.py           | 18 +++++++++
 .../healthcheck/healthcheck_analysis.py       | 38 +++++++++++++++++++
 2 files changed, 56 insertions(+)
 create mode 100644 ax/analysis/healthcheck/__init__.py
 create mode 100644 ax/analysis/healthcheck/healthcheck_analysis.py

diff --git a/ax/analysis/healthcheck/__init__.py b/ax/analysis/healthcheck/__init__.py
new file mode 100644
index 00000000000..2d99e3e9a0d
--- /dev/null
+++ b/ax/analysis/healthcheck/__init__.py
@@ -0,0 +1,18 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-strict
+
+from ax.analysis.healthcheck.healthcheck_analysis import (
+    HealthcheckAnalysis,
+    HealthcheckAnalysisCard,
+    HealthcheckStatus,
+)
+
+__all__ = [
+    "HealthcheckAnalysis",
+    "HealthcheckAnalysisCard",
+    "HealthcheckStatus",
+]
diff --git a/ax/analysis/healthcheck/healthcheck_analysis.py b/ax/analysis/healthcheck/healthcheck_analysis.py
new file mode 100644
index 00000000000..c7fd74d025e
--- /dev/null
+++ b/ax/analysis/healthcheck/healthcheck_analysis.py
@@ -0,0 +1,38 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-strict
+import json
+from enum import IntEnum
+from typing import Optional
+
+from ax.analysis.analysis import AnalysisCard
+from ax.core.experiment import Experiment
+from ax.modelbridge.generation_strategy import GenerationStrategy
+
+
+class HealthcheckStatus(IntEnum):
+    PASS = 0
+    FAIL = 1
+    WARNING = 2
+
+
+class HealthcheckAnalysisCard(AnalysisCard):
+    blob_annotation = "healthcheck"
+
+    def get_status(self) -> HealthcheckStatus:
+        return HealthcheckStatus(json.loads(self.blob)["status"])
+
+
+class HealthcheckAnalysis:
+    """
+    An analysis that performs a health check.
+    """
+
+    def compute(
+        self,
+        experiment: Optional[Experiment] = None,
+        generation_strategy: Optional[GenerationStrategy] = None,
+    ) -> HealthcheckAnalysisCard: ...