Register NotApplicableStateAnalysisCard in JSON store registry (#5163)

Summary:
Pull Request resolved: #5163

Register `NotApplicableStateAnalysisCard` in the JSON store encoder/decoder registries for completeness and consistency with other AnalysisCard subclasses.

The JSON store registry uses exact type matching, so every concrete AnalysisCard subclass needs its own entry. This was not needed until now because analysis cards are persisted exclusively through the SQA store (which uses `isinstance()` checks and does not rely on this registry).
Adding the registration ensures the JSON store can correctly round-trip `NotApplicableStateAnalysisCard`, keeping it consistent with sibling subclasses like `ErrorAnalysisCard` and guarding against future use cases that may serialize analysis cards via JSON.

Reviewed By: Cesar-Cardoso

Differential Revision: D100189187

fbshipit-source-id: 44374b22d3b088ac70b7c5fb2833726009d92023

Author: shrutipatel31

ax/storage/json_store/registry.py

@@ -28,7 +28,12 @@
 from ax.benchmark.benchmark_result import AggregatedBenchmarkResult, BenchmarkResult
 from ax.benchmark.benchmark_trial_metadata import BenchmarkTrialMetadata
 from ax.core import Experiment, ObservationFeatures
-from ax.core.analysis_card import AnalysisCard, AnalysisCardGroup, ErrorAnalysisCard
+from ax.core.analysis_card import (
+    AnalysisCard,
+    AnalysisCardGroup,
+    ErrorAnalysisCard,
+    NotApplicableStateAnalysisCard,
+)
 from ax.core.arm import Arm
 from ax.core.auxiliary import AuxiliaryExperiment, AuxiliaryExperimentPurpose
 from ax.core.batch_trial import AbandonedArm, BatchTrial
@@ -220,6 +225,7 @@
     ExpressionDerivedMetric: metric_to_dict,
     DerivedParameter: derived_parameter_to_dict,
     ErrorAnalysisCard: analysis_card_to_dict,
+    NotApplicableStateAnalysisCard: analysis_card_to_dict,
     Experiment: experiment_to_dict,
     FactorialMetric: metric_to_dict,
     FixedParameter: fixed_parameter_to_dict,
@@ -345,6 +351,7 @@
     "DerivedParameter": DerivedParameter,
     "DomainType": DomainType,
     "ErrorAnalysisCard": ErrorAnalysisCard,
+    "NotApplicableStateAnalysisCard": NotApplicableStateAnalysisCard,
     "Experiment": Experiment,
     "ExperimentStatus": ExperimentStatus,
     "FactorialMetric": FactorialMetric,

ax/storage/json_store/tests/test_json_store.py

@@ -37,7 +37,12 @@
     get_benchmark_result,
     get_benchmark_time_varying_metric,
 )
-from ax.core.analysis_card import AnalysisCard, AnalysisCardGroup, ErrorAnalysisCard
+from ax.core.analysis_card import (
+    AnalysisCard,
+    AnalysisCardGroup,
+    ErrorAnalysisCard,
+    NotApplicableStateAnalysisCard,
+)
 from ax.core.auxiliary import AuxiliaryExperimentPurpose
 from ax.core.data import Data
 from ax.core.generator_run import GeneratorRun
@@ -473,6 +478,16 @@
             blob="digraph {}",
         ),
     ),
+    (
+        "NotApplicableStateAnalysisCard",
+        lambda: NotApplicableStateAnalysisCard(
+            name="TestNA",
+            title="Not Applicable",
+            subtitle="na subtitle",
+            df=pd.DataFrame(),
+            blob="Not enough data.",
+        ),
+    ),
     (
         "AnalysisCardGroup",
         lambda: AnalysisCardGroup(

ax/storage/sqa_store/tests/test_sqa_store.py

@@ -24,7 +24,11 @@
     choose_generation_strategy,
     GenerationStrategyDispatchStruct,
 )
-from ax.core.analysis_card import AnalysisCard, AnalysisCardGroup
+from ax.core.analysis_card import (
+    AnalysisCard,
+    AnalysisCardGroup,
+    NotApplicableStateAnalysisCard,
+)
 from ax.core.arm import Arm
 from ax.core.auxiliary import (
     AuxiliaryExperiment,
@@ -3030,14 +3034,26 @@ def test_analysis_card(self) -> None:
             df=test_df,
             blob=pio.to_json(go.Figure()),
         )
+        na_card = NotApplicableStateAnalysisCard(
+            name="test_na_card",
+            title="test_title",
+            subtitle="test_subtitle",
+            df=test_df,
+            blob="Not enough data.",
+        )
 
         # Create two groups which hold the leaf cards
         # Add the same analysis card multiple times to test _unique_id logic
         small_group = AnalysisCardGroup(
             name="small_group",
             title="Small Group",
             subtitle="This is a small group with just a few cards",
-            children=[base_analysis_card, markdown_analysis_card, plotly_analysis_card],
+            children=[
+                base_analysis_card,
+                markdown_analysis_card,
+                plotly_analysis_card,
+                na_card,
+            ],
         )
         big_group = AnalysisCardGroup(
             name="big_group",
@@ -3099,6 +3115,12 @@ def test_analysis_card(self) -> None:
             self.assertEqual(loaded_small_group_plotly.name, plotly_analysis_card.name)
             self.assertEqual(loaded_small_group_plotly.blob, plotly_analysis_card.blob)
 
+            loaded_na = assert_is_instance(
+                loaded_small_group.children[3], NotApplicableStateAnalysisCard
+            )
+            self.assertEqual(loaded_na.name, na_card.name)
+            self.assertEqual(loaded_na.blob, na_card.blob)
+
     def test_delete_generation_strategy(self) -> None:
         # GIVEN an experiment with a generation strategy
         experiment = get_branin_experiment()