fix: improve enable_data_report handling

mostlyai/sdk/_local/execution/plan.py

@@ -48,13 +48,6 @@
 ]
 
 
-def _get_keys(table: SourceTable) -> list[str]:
-    keys = [fk.column for fk in (table.foreign_keys or [])]
-    if table.primary_key:
-        keys.append(table.primary_key)
-    return keys
-
-
 def has_tabular_model(table: SourceTable) -> bool:
     return table.tabular_model_configuration is not None
 
@@ -63,13 +56,22 @@ def has_language_model(table: SourceTable) -> bool:
     return table.language_model_configuration is not None
 
 
-def get_model_type_generation_steps_map(include_report: bool) -> dict[ModelType, list[StepCode]]:
-    return {
-        ModelType.tabular: [StepCode.generate_data_tabular]
-        + ([StepCode.create_data_report_tabular] if include_report else []),
-        ModelType.language: [StepCode.generate_data_language]
-        + ([StepCode.create_data_report_language] if include_report else []),
+def get_model_type_generation_steps_map(
+    enable_data_report: bool, table: SourceTable
+) -> dict[ModelType, list[StepCode]]:
+    steps = {
+        ModelType.tabular: [StepCode.generate_data_tabular],
+        ModelType.language: [StepCode.generate_data_language],
     }
+    if enable_data_report and has_tabular_model(table):
+        enable_tabular_model_report = table.tabular_model_configuration.enable_model_report
+        if enable_tabular_model_report:
+            steps[ModelType.tabular].append(StepCode.create_data_report_tabular)
+    if enable_data_report and has_language_model(table):
+        enable_language_model_report = table.language_model_configuration.enable_model_report
+        if enable_language_model_report:
+            steps[ModelType.language].append(StepCode.create_data_report_language)
+    return steps
 
 
 class Step(BaseModel):
@@ -159,7 +161,7 @@ def make_generator_execution_plan(generator: Generator) -> ExecutionPlan:
 
 
 def make_synthetic_dataset_execution_plan(
-    generator: Generator, synthetic_dataset: SyntheticDataset | None = None, is_probe: bool = False
+    generator: Generator, synthetic_dataset: SyntheticDataset, is_probe: bool = False
 ) -> ExecutionPlan:
     execution_plan = ExecutionPlan(tasks=[])
     sync_task = execution_plan.add_task(TaskType.sync)
@@ -168,20 +170,19 @@ def make_synthetic_dataset_execution_plan(
     generate_steps = []
 
     def add_generation_steps(table: SourceTable):
-        if synthetic_dataset:
-            synthetic_table = next(t for t in synthetic_dataset.tables if t.name == table.name)
-            enable_data_report = synthetic_table.configuration.enable_data_report
-        else:
-            enable_data_report = True
+        synthetic_table = next(t for t in synthetic_dataset.tables if t.name == table.name)
+        enable_data_report = synthetic_table.configuration.enable_data_report
         if has_tabular_model(table):
+            enable_tabular_model_report = table.tabular_model_configuration.enable_model_report
             steps = [Step(step_code=StepCode.generate_data_tabular, target_table_name=table.name)]
-            if not is_probe and enable_data_report:
+            if not is_probe and enable_data_report and enable_tabular_model_report:
                 steps.append(Step(step_code=StepCode.create_data_report_tabular, target_table_name=table.name))
             generate_steps.extend(steps)
 
         if has_language_model(table):
+            enable_language_model_report = table.language_model_configuration.enable_model_report
             steps = [Step(step_code=StepCode.generate_data_language, target_table_name=table.name)]
-            if not is_probe and enable_data_report:
+            if not is_probe and enable_data_report and enable_language_model_report:
                 steps.append(Step(step_code=StepCode.create_data_report_language, target_table_name=table.name))
             generate_steps.extend(steps)
 

mostlyai/sdk/_local/synthetic_datasets.py

@@ -119,7 +119,10 @@ def create_synthetic_dataset(
     progress_steps: list[ProgressStep] = []
     for table in generator.tables:
         sd_table = next(t for t in config.tables if t.name == table.name)
-        steps_map = get_model_type_generation_steps_map(sd_table.configuration.enable_data_report)
+        steps_map = get_model_type_generation_steps_map(
+            enable_data_report=sd_table.configuration.enable_data_report,
+            table=table,
+        )
         model_types = [
             model_type
             for model_type, check in [

mostlyai/sdk/domain.py

@@ -3643,7 +3643,7 @@ def validate_data_report_disabled_if_both_model_reports_disabled(cls, validation
             ]
             if cfg
         ]
-        if not all(cfg.enable_model_report for cfg in configs):
+        if all(cfg.enable_model_report is False for cfg in configs):
             if validation.synthetic_table.configuration is not None:
                 validation.synthetic_table.configuration.enable_data_report = False
         return validation

tests/_local/end_to_end/test_simple_flat.py

@@ -207,7 +207,7 @@ def test_reproducibility(tmp_path):
     )
     g_config = {
         "random_state": 42,
-        "tables": [{"name": "data", "data": df, "model_configuration": {"max_epochs": 0.1}}],
+        "tables": [{"name": "data", "data": df, "tabular_model_configuration": {"max_epochs": 0.1}}],
     }
     g1 = mostly.train(config=g_config)
     g2 = mostly.train(config=g_config)

tests/_local/unit/test_execution_plan.py

@@ -91,6 +91,7 @@ def test_make_synthetic_dataset_execution_plan():
                         {"name": "order_id", "model_encoding_type": "AUTO"},
                     ],
                     "foreign_keys": [{"column": "order_id", "referenced_table": "orders", "is_context": True}],
+                    "tabular_model_configuration": {"enable_model_report": False},
                 },
             ],
         }
@@ -123,7 +124,7 @@ def test_make_synthetic_dataset_execution_plan():
             Step(step_code=StepCode.generate_data_tabular, target_table_name="orders"),
             # No data report for orders table
             Step(step_code=StepCode.generate_data_tabular, target_table_name="order_items"),
-            Step(step_code=StepCode.create_data_report_tabular, target_table_name="order_items"),
+            # No tabular data report for order_items table
             Step(step_code=StepCode.generate_data_language, target_table_name="order_items"),
             Step(step_code=StepCode.create_data_report_language, target_table_name="order_items"),
             Step(step_code=StepCode.generate_data_tabular, target_table_name="prices"),
@@ -223,7 +224,7 @@ def test_make_generator_execution_plan():
 
 def test_make_synthetic_dataset_execution_plan_with_probe():
     model_config = arbitrary_model_config()
-    config = Generator(
+    generator = Generator(
         name="test_probe_dataset",
         training_status=ProgressStatus.in_progress,
         tables=[
@@ -258,8 +259,13 @@ def test_make_synthetic_dataset_execution_plan_with_probe():
             ),
         ],
     )
+    synthetic_dataset = SyntheticDataset(
+        tables=[
+            SyntheticTable(name=table.name, configuration=SyntheticTableConfiguration()) for table in generator.tables
+        ]
+    )
 
-    execution_plan = make_synthetic_dataset_execution_plan(config, is_probe=True)
+    execution_plan = make_synthetic_dataset_execution_plan(generator, synthetic_dataset, is_probe=True)
 
     expected_execution_plan = ExecutionPlan(tasks=[])
     sync_task = expected_execution_plan.add_task(TaskType.sync)