Improve error message when using the distributed gbt learner with a n…

…on supported task. PiperOrigin-RevId: 673362728
google · Sep 11, 2024 · b7f5472 · b7f5472
1 parent 3e22c90
commit b7f5472
Showing 1 changed file with 12 additions and 2 deletions.
diff --git a/yggdrasil_decision_forests/learner/distributed_gradient_boosted_trees/worker.cc b/yggdrasil_decision_forests/learner/distributed_gradient_boosted_trees/worker.cc
@@ -434,7 +434,12 @@ absl::Status DistributedGradientBoostedTreesWorker::GetLabelStatistics(
           answer->mutable_label_statistics(), thread_pool_.get()));
     } break;
     default:
-      return absl::InvalidArgumentError("Not supported task");
+      return absl::InvalidArgumentError(
+          absl::StrCat("The distributed gradient boosted trees learner does "
+                       "not support this task:",
+                       model::proto::Task_Name(welcome_.train_config().task()),
+                       ". Use the non-distributed "
+                       "version of the learner."));
   }
 
   return absl::OkStatus();
@@ -1267,7 +1272,12 @@ absl::StatusOr<LossResults> DistributedGradientBoostedTreesWorker::Loss(
       break;
     }
     default:
-      return absl::InvalidArgumentError("Not supported task");
+      return absl::InvalidArgumentError(
+          absl::StrCat("The distributed gradient boosted trees learner does "
+                       "not support this task:",
+                       model::proto::Task_Name(welcome_.train_config().task()),
+                       ". Use the non-distributed "
+                       "version of the learner."));
   }
 }