Added Plexus as an Airflow provider (#10591)

CONTRIBUTING.rst

@@ -502,7 +502,7 @@ celery, cgroups, cloudant, cncf.kubernetes, dask, databricks, datadog, devel, de
 docker, druid, elasticsearch, exasol, facebook, gcp, gcp_api, github_enterprise, google,
 google_auth, grpc, hashicorp, hdfs, hive, jdbc, jira, kerberos, kubernetes, ldap, microsoft.azure,
 microsoft.mssql, microsoft.winrm, mongo, mssql, mysql, odbc, oracle, pagerduty, papermill, password,
-pinot, postgres, presto, qds, rabbitmq, redis, salesforce, samba, segment, sendgrid, sentry,
+pinot, plexus, postgres, presto, qds, rabbitmq, redis, salesforce, samba, segment, sendgrid, sentry,
 singularity, slack, snowflake, spark, ssh, statsd, tableau, vertica, virtualenv, webhdfs, winrm,
 yandexcloud, all, devel_ci
 

INSTALL

@@ -51,7 +51,7 @@ celery, cgroups, cloudant, cncf.kubernetes, dask, databricks, datadog, devel, de
 docker, druid, elasticsearch, exasol, facebook, gcp, gcp_api, github_enterprise, google,
 google_auth, grpc, hashicorp, hdfs, hive, jdbc, jira, kerberos, kubernetes, ldap, microsoft.azure,
 microsoft.mssql, microsoft.winrm, mongo, mssql, mysql, odbc, oracle, pagerduty, papermill, password,
-pinot, postgres, presto, qds, rabbitmq, redis, salesforce, samba, segment, sendgrid, sentry,
+pinot, plexus, postgres, presto, qds, rabbitmq, redis, salesforce, samba, segment, sendgrid, sentry,
 singularity, slack, snowflake, spark, ssh, statsd, tableau, vertica, virtualenv, webhdfs, winrm,
 yandexcloud, all, devel_ci
 

airflow/providers/plexus/__init__.py

@@ -0,0 +1,16 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.

airflow/providers/plexus/example_dags/__init__.py

@@ -0,0 +1,16 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.

airflow/providers/plexus/example_dags/example_plexus.py

@@ -0,0 +1,49 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from airflow import DAG
+from airflow.providers.plexus.operators.job import PlexusJobOperator
+from airflow.utils.dates import days_ago
+
+HOME = '/home/acc'
+T3_PRERUN_SCRIPT = 'cp {home}/imdb/run_scripts/mlflow.sh {home}/ && chmod +x mlflow.sh'.format(home=HOME)
+
+args = {'owner': 'core scientific', 'retries': 1}
+
+dag = DAG(
+    'test',
+    default_args=args,
+    description='testing plexus operator',
+    start_date=days_ago(1),
+    schedule_interval='@once',
+    catchup=False,
+)
+
+t1 = PlexusJobOperator(
+    task_id='test',
+    job_params={
+        'name': 'test',
+        'app': 'MLFlow Pipeline 01',
+        'queue': 'DGX-2 (gpu:Tesla V100-SXM3-32GB)',
+        'num_nodes': 1,
+        'num_cores': 1,
+        'prerun_script': T3_PRERUN_SCRIPT,
+    },
+    dag=dag,
+)
+
+t1

airflow/providers/plexus/hooks/__init__.py

@@ -0,0 +1,16 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.

airflow/providers/plexus/hooks/plexus.py

@@ -0,0 +1,73 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import jwt
+import arrow
+import requests
+from airflow.hooks.base_hook import BaseHook
+from airflow.models import Variable
+from airflow.exceptions import AirflowException
+
+
+class PlexusHook(BaseHook):
+    """
+    Used for jwt token generation and storage to
+    make Plexus API calls. Requires email and password
+    Airflow variables be created.
+
+    Example:
+        - export AIRFLOW_VAR_EMAIL = user@corescientific.com
+        - export AIRFLOW_VAR_PASSWORD = *******
+
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.__token = None
+        self.__token_exp = None
+        self.host = "https://apiplexus.corescientific.com/"
+        self.user_id = None
+
+    def _generate_token(self):
+        login = Variable.get("email")
+        pwd = Variable.get("password")
+        if login is None or pwd is None:
+            raise AirflowException("No valid email/password supplied.")
+        token_endpoint = self.host + "sso/jwt-token/"
+        response = requests.post(token_endpoint, data={"email": login, "password": pwd}, timeout=5)
+        if not response.ok:
+            raise AirflowException(
+                "Could not retrieve JWT Token. Status Code: [{}]. "
+                "Reason: {} - {}".format(response.status_code, response.reason, response.text)
+            )
+        token = response.json()["access"]
+        payload = jwt.decode(token, verify=False)
+        self.user_id = payload["user_id"]
+        self.__token_exp = payload["exp"]
+
+        return token
+
+    @property
+    def token(self):
+        """Returns users token"""
+        if self.__token is not None:
+            if arrow.get(self.__token_exp) <= arrow.now():
+                self.__token = self._generate_token()
+            return self.__token
+        else:
+            self.__token = self._generate_token()
+            return self.__token

airflow/providers/plexus/operators/__init__.py

@@ -0,0 +1,16 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.

airflow/providers/plexus/operators/job.py

@@ -0,0 +1,162 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import time
+import logging
+from typing import Dict
+import requests
+from airflow.providers.plexus.hooks.plexus import PlexusHook
+from airflow.models import BaseOperator
+from airflow.utils.decorators import apply_defaults
+from airflow.exceptions import AirflowException
+
+logger = logging.getLogger(__name__)
+
+
+class PlexusJobOperator(BaseOperator):
+    """
+    Submits a Plexus job.
+
+    :param job_params: parameters required to launch a job.
+    :type job_params: dict
+
+    Required job parameters are the following
+        - "name": job name created by user.
+        - "app": name of the application to run. found in Plexus UI.
+        - "queue": public cluster name. found in Plexus UI.
+        - "num_nodes": number of nodes.
+        - "num_cores":  number of cores per node.
+
+    """
+
+    @apply_defaults
+    def __init__(self, job_params: Dict, **kwargs) -> None:
+        super().__init__(**kwargs)
+
+        self.job_params = job_params
+        self.required_params = set(["name", "app", "queue", "num_cores", "num_nodes"])
+        self.lookups = {
+            "app": ("apps/", "id", "name"),
+            "billing_account_id": ("users/{}/billingaccounts/", "id", None),
+            "queue": ("queues/", "id", "public_name"),
+        }
+        self.job_params.update({"billing_account_id": None})
+        self.is_service = None
+
+    def execute(self, context):
+        hook = PlexusHook()
+        params = self.construct_job_params(hook)
+        if self.is_service is True:
+            if self.job_params.get("expected_runtime") is None:
+                end_state = "Running"
+            else:
+                end_state = "Finished"
+        elif self.is_service is False:
+            end_state = "Completed"
+        else:
+            raise AirflowException(
+                "Unable to determine if application "
+                "is running as a batch job or service. "
+                "Contact Core Scientific AI Team."
+            )
+        logger.info("creating job w/ following params: %s", params)
+        jobs_endpoint = hook.host + "jobs/"
+        headers = {"Authorization": "Bearer {}".format(hook.token)}
+        create_job = requests.post(jobs_endpoint, headers=headers, data=params, timeout=5)
+        if create_job.ok:
+            job = create_job.json()
+            jid = job["id"]
+            state = job["last_state"]
+            while state != end_state:
+                time.sleep(3)
+                jid_endpoint = jobs_endpoint + "{}/".format(jid)
+                get_job = requests.get(jid_endpoint, headers=headers, timeout=5)
+                if not get_job.ok:
+                    raise AirflowException(
+                        "Could not retrieve job status. Status Code: [{0}]. "
+                        "Reason: {1} - {2}".format(get_job.status_code, get_job.reason, get_job.text)
+                    )
+                new_state = get_job.json()["last_state"]
+                if new_state in ("Cancelled", "Failed"):
+                    raise AirflowException("Job {}".format(new_state))
+                elif new_state != state:
+                    logger.info("job is %s", new_state)
+                state = new_state
+        else:
+            raise AirflowException(
+                "Could not start job. Status Code: [{}]. "
+                "Reason: {} - {}".format(create_job.status_code, create_job.reason, create_job.text)
+            )
+
+    def _api_lookup(self, param: str, hook):
+        lookup = self.lookups[param]
+        key = lookup[1]
+        mapping = None if lookup[2] is None else (lookup[2], self.job_params[param])
+
+        if param == "billing_account_id":
+            endpoint = hook.host + lookup[0].format(hook.user_id)
+        else:
+            endpoint = hook.host + lookup[0]
+        headers = {"Authorization": "Bearer {}".format(hook.token)}
+        response = requests.get(endpoint, headers=headers, timeout=5)
+        results = response.json()["results"]
+
+        v = None
+        if mapping is None:
+            v = results[0][key]
+        else:
+            for dct in results:
+                if dct[mapping[0]] == mapping[1]:
+                    v = dct[key]
+                if param == 'app':
+                    self.is_service = dct['is_service']
+        if v is None:
+            raise AirflowException(
+                "Could not locate value for param:{} at endpoint: {}".format(key, endpoint)
+            )
+
+        return v
+
+    def construct_job_params(self, hook):
+        """
+        Creates job_params dict for api call to
+        launch a Plexus job.
+
+        Some parameters required to launch a job
+        are not available to the user in the Plexus
+        UI. For example, an app id is required, but
+        only the app name is provided in the UI.
+        This function acts as a backend lookup
+        of the required param value using the
+        user-provided value.
+
+        :param hook: plexus hook object
+        :type hook: airflow hook
+        """
+        missing_params = self.required_params - set(self.job_params)
+        if len(missing_params) > 0:
+            raise AirflowException(
+                "Missing the following required job_params: {}".format(", ".join(missing_params))
+            )
+        params = {}
+        for prm in self.job_params:
+            if prm in self.lookups:
+                v = self._api_lookup(param=prm, hook=hook)
+                params[prm] = v
+            else:
+                params[prm] = self.job_params[prm]
+        return params

docs/autoapi_templates/index.rst

@@ -180,6 +180,8 @@ All operators are in the following packages:
 
   airflow/providers/papermill/operators/index
 
+  airflow/providers/plexus/operators/index
+
   airflow/providers/postgres/operators/index
 
   airflow/providers/qubole/operators/index
@@ -325,6 +327,8 @@ All hooks are in the following packages:
 
   airflow/providers/pagerduty/hooks/index
 
+  airflow/providers/plexus/hooks/index
+
   airflow/providers/postgres/hooks/index
 
   airflow/providers/presto/hooks/index