Add on-demand 'trigger-ml' CLI command for machine learning workers (fixes #2920)

augur/application/cli/tasks.py

@@ -4,17 +4,18 @@
 Augur library commands for controlling the backend components
 """
 
 import os
 import click
 import subprocess
 import uuid
 import time
 # from redis.exceptions import ConnectionError as RedisConnectionError
 
 from augur import instance_id
 from augur.application.logs import AugurLogger
 from augur.application.cli import test_connection, test_db_connection 
 from augur.application.cli.backend import clear_rabbitmq_messages, raise_open_file_limit
+from augur.tasks.start_tasks import trigger_ml_phase
 
 logger = AugurLogger("augur", reset_logfiles=False).get_logger()
 
@@ -60,7 +61,7 @@
             secondary_worker_process.terminate()
 
         try:
             clear_rabbitmq_messages()
 
         except Exception as e:
             pass
@@ -91,4 +92,13 @@
         else:
             logger.error("Invalid input")
 
-
+
+@cli.command("trigger-ml")
+@test_connection
+@test_db_connection
+def trigger_ml():
+    logger.info("Starting ML tasks")
+    trigger_ml_phase.delay()
+    logger.info("ML tasks started")
+
+

augur/tasks/start_tasks.py

@@ -236,6 +236,27 @@
     request.get_valid_repos(session)
     return request
 
+@celery.task(bind=True)
+def trigger_ml_phase(self):
+    engine = self.app.engine
+    logger = logging.getLogger(trigger_ml_phase.__name__)
+    enabled_phase_names = get_enabled_phase_names_from_config(engine, logger)
+
+    if RUNNING_DOCKER or machine_learning_phase.__name__ not in enabled_phase_names:
+        return 
+
+
+    with DatabaseSession(logger, engine) as session:
+        config = AugurConfig(logger, session)
+        ml_phase_request = []
+        ml_interval = config.get_value('Tasks', 'ml_collection_interval_days') or 40
+        ml_phase_request.append(build_ml_repo_collect_request(session, logger, enabled_phase_names, ml_interval))
+
+    main_routine = AugurTaskRoutine(logger, ml_phase_request)
+
+    main_routine.start_data_collection()
+
+
 @celery.task(bind=True)
 def augur_collection_monitor(self):     
 
@@ -271,10 +292,6 @@
             #start_facade_collection(session, max_repo=30)
             enabled_collection_hooks.append(build_facade_repo_collect_request(session, logger, enabled_phase_names, facade_interval))
 
-        if not RUNNING_DOCKER and machine_learning_phase.__name__ in enabled_phase_names:
-            enabled_collection_hooks.append(build_ml_repo_collect_request(session, logger, enabled_phase_names, ml_interval))
-            #start_ml_collection(session,max_repo=5)
-
         logger.info(f"Starting collection phases: {[h.name for h in enabled_collection_hooks]}")
 
         main_routine = AugurTaskRoutine(logger, enabled_collection_hooks)
@@ -304,7 +321,7 @@
             status = repo.collection_status[0]
             raw_count = status.issue_pr_sum
 
             issue_pr_task_update_weight_util([int(raw_count)],repo_git=repo_git,session=session)
 
         facade_not_pending = CollectionStatus.facade_status != CollectionState.PENDING.value
         facade_not_failed = CollectionStatus.facade_status != CollectionState.FAILED_CLONE.value