fix python segfault

flexflow · Jul 18, 2023 · 716a5b1 · 716a5b1
1 parent f7ad4ba
commit 716a5b1
Showing 1 changed file with 54 additions and 15 deletions.
diff --git a/src/runtime/model.cc b/src/runtime/model.cc
@@ -4103,20 +4103,40 @@ void register_flexflow_internal_tasks(Runtime *runtime,
                                    "RequestManager Prepare Next Batch");
     registrar.add_constraint(ProcessorConstraint(Processor::LOC_PROC));
     registrar.set_leaf();
-    Runtime::preregister_task_variant<BatchConfig,
-                                      RequestManager::prepare_next_batch_task>(
-        registrar, "RequestManager Prepare Next Batch Task");
+    if (pre_register) {
+      Runtime::preregister_task_variant<
+          BatchConfig,
+          RequestManager::prepare_next_batch_task>(
+          registrar, "RequestManager Prepare Next Batch Task");
+    } else {
+      if (enable_control_replication) {
+        registrar.global_registration = false;
+      }
+      runtime->register_task_variant<BatchConfig,
+                                     RequestManager::prepare_next_batch_task>(
+          registrar);
+    }
   }
   // RequestManager prepare_next_batch_beam
   {
     TaskVariantRegistrar registrar(RM_PREPARE_NEXT_BATCH_BEAM_TASK_ID,
                                    "RequestManager Prepare Next Batch (Beam)");
     registrar.add_constraint(ProcessorConstraint(Processor::LOC_PROC));
     registrar.set_leaf();
-    Runtime::preregister_task_variant<
-        BeamSearchBatchConfig,
-        RequestManager::prepare_next_batch_beam_task>(
-        registrar, "RequestManager Prepare Next Batch (Beam) Task");
+    if (pre_register) {
+      Runtime::preregister_task_variant<
+          BeamSearchBatchConfig,
+          RequestManager::prepare_next_batch_beam_task>(
+          registrar, "RequestManager Prepare Next Batch (Beam) Task");
+    } else {
+      if (enable_control_replication) {
+        registrar.global_registration = false;
+      }
+      runtime
+          ->register_task_variant<BeamSearchBatchConfig,
+                                  RequestManager::prepare_next_batch_beam_task>(
+              registrar);
+    }
   }
   // RequestManager prepare_next_batch_init
   {
@@ -4125,10 +4145,20 @@ void register_flexflow_internal_tasks(Runtime *runtime,
         "RequestManager Prepare Next Batch (Init Beam)");
     registrar.add_constraint(ProcessorConstraint(Processor::LOC_PROC));
     registrar.set_leaf();
-    Runtime::preregister_task_variant<
-        BeamSearchBatchConfig,
-        RequestManager::prepare_next_batch_init_task>(
-        registrar, "RequestManager Prepare Next Batch (Init Beam) Task");
+    if (pre_register) {
+      Runtime::preregister_task_variant<
+          BeamSearchBatchConfig,
+          RequestManager::prepare_next_batch_init_task>(
+          registrar, "RequestManager Prepare Next Batch (Init Beam) Task");
+    } else {
+      if (enable_control_replication) {
+        registrar.global_registration = false;
+      }
+      runtime
+          ->register_task_variant<BeamSearchBatchConfig,
+                                  RequestManager::prepare_next_batch_init_task>(
+              registrar);
+    }
   }
   // RequestManager prepare_next_batch_verify
   {
@@ -4137,10 +4167,19 @@ void register_flexflow_internal_tasks(Runtime *runtime,
         "RequestManager Prepare Next Batch (Verify)");
     registrar.add_constraint(ProcessorConstraint(Processor::LOC_PROC));
     registrar.set_leaf();
-    Runtime::preregister_task_variant<
-        TreeVerifyBatchConfig,
-        RequestManager::prepare_next_batch_verify_task>(
-        registrar, "RequestManager Prepare Next Batch (Verify) Task");
+    if (pre_register) {
+      Runtime::preregister_task_variant<
+          TreeVerifyBatchConfig,
+          RequestManager::prepare_next_batch_verify_task>(
+          registrar, "RequestManager Prepare Next Batch (Verify) Task");
+    } else {
+      if (enable_control_replication) {
+        registrar.global_registration = false;
+      }
+      runtime->register_task_variant<
+          TreeVerifyBatchConfig,
+          RequestManager::prepare_next_batch_verify_task>(registrar);
+    }
   }
   // ElementUnary task
   {