Merge branch 'add-support-for-custom-scheduler' into 'master'

Nicanor Gutiérrez Requejo · Nicanor Gutiérrez Requejo · commit e4470da4ff03 · 2021-07-14T14:58:13.000Z
Add support for custom scheduler

See merge request wuxi-nextcode/research-platform/backend-services/platform/python-sdk!52
diff --git a/nextcode/VERSION b/nextcode/VERSION
@@ -1 +1 @@
-1.14.1
+1.15.0
diff --git a/nextcode/services/workflow/pytest_plugin.py b/nextcode/services/workflow/pytest_plugin.py
@@ -95,6 +95,13 @@ def pytest_addoption(parser):
         help="Specifies the credentials to forward to workflow-service",
         default=[],
     )
+    group.addoption(
+        "--scheduler_name",
+        action="store",
+        dest="scheduler_name",
+        help="Specifies the scheduler to use in the pipeline run",
+        default=None,
+    )
     parser.addini(
         "base_upload_bucket",
         "The S3 location the result_dir parameter of the pipeline will be set to",
@@ -166,6 +173,7 @@ def run_workflow(request):
     run_id = request.config.option.run_id
     storage_type = request.config.option.storage_type
     credentials_args = request.config.option.credentials
+    scheduler_name = request.config.option.scheduler_name
     # Use the profile specified by the TestCase if it exists. Use the default profile if none is provided.
     profile = getattr(request.cls, "profile", request.config.option.profile)
     revision = None
@@ -215,6 +223,7 @@ def run_workflow(request):
         profile=profile,
         storage_type=storage_type,
         credentials=credentials,
+        scheduler_name=scheduler_name,
     )
 
     last_status = ""
diff --git a/nextcode/services/workflow/service.py b/nextcode/services/workflow/service.py
@@ -187,6 +187,7 @@ def post_job(
         keep_resources: Optional[bool] = None,
         mount_jobs: Optional[List[int]] = None,
         credentials: Optional[Dict] = None,
+        scheduler_name: Optional[str] = None,
     ):
         """
         Run a workflow job
@@ -224,6 +225,8 @@ def post_job(
                 'aws_secret_access_key': 'k6q5...'
             }
         }
+        :param scheduler_name: Optional string specifying the scheduler to use. If not specified, the default scheduler
+               will be used
         """
         if not project_name:
             project_name = os.environ.get("GOR_API_PROJECT")
@@ -234,7 +237,7 @@ def post_job(
         log.debug(
             "post_job called with pipeline_name=%s, project_name=%s, params=%s, script=%s, revision=%s, "
             "build_source=%s, build_context=%s, profile=%s, description=%s, executor_memory_mb=%s, context=%s, "
-            "storage_type=%s, dedicated_storage_size=%s, mount_jobs=%s",
+            "storage_type=%s, dedicated_storage_size=%s, mount_jobs=%s scheduler_name=%s",
             pipeline_name,
             project_name,
             params,
@@ -249,6 +252,7 @@ def post_job(
             storage_type,
             dedicated_storage_size,
             mount_jobs,
+            scheduler_name,
         )
         data: Dict = {
             "pipeline_name": pipeline_name,
@@ -278,6 +282,8 @@ def post_job(
             data["keep_resources"] = keep_resources
         if mount_jobs:
             data["mount_jobs"] = mount_jobs
+        if scheduler_name:
+            data["scheduler_name"] = scheduler_name
 
         endpoint = self.session.url_from_endpoint("jobs")
 
diff --git a/tests/test_workflow.py b/tests/test_workflow.py
@@ -247,6 +247,7 @@ def test_post_job(self):
         build_context = None
         profile = "test"
         storage_type = "dedicated"
+        scheduler_name = "custom-scheduler"
 
         job = self.svc.post_job(
             pipeline_name,
@@ -258,6 +259,7 @@ def test_post_job(self):
             build_context,
             profile,
             storage_type,
+            scheduler_name,
         )
         self.assertEqual(666, job.job_id)
 
@@ -280,6 +282,7 @@ def test_post_job_with_creds(self):
                 "aws_secret_access_key": "Aj39sadfljhdslafjasls",
             }
         }
+        scheduler_name = "custom-scheduler"
 
         job = self.svc.post_job(
             pipeline_name,
@@ -292,6 +295,7 @@ def test_post_job_with_creds(self):
             profile,
             storage_type,
             credentials,
+            scheduler_name,
         )
         self.assertEqual(666, job.job_id)
 
