Migrate Plexus example DAGs to new design apache#22457

chethanuk · Jun 3, 2022 · 0e11fa3 · 0e11fa3
1 parent f294a26
commit 0e11fa3
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 33 deletions.
diff --git a/airflow/providers/plexus/example_dags/__init__.py b/airflow/providers/plexus/example_dags/__init__.py
diff --git a/docs/apache-airflow-providers-plexus/index.rst b/docs/apache-airflow-providers-plexus/index.rst
@@ -32,7 +32,7 @@ Content
     :maxdepth: 1
     :caption: Resources
 
-    Example DAGs <https://github.com/apache/airflow/tree/main/airflow/providers/plexus/example_dags>
+    Example DAGs <https://github.com/apache/airflow/tree/main/tests/system/providers/plexus>
 
 .. toctree::
     :maxdepth: 1

diff --git a/...ers/plexus/example_dags/example_plexus.py → ...system/providers/plexus/example_plexus.py b/...ers/plexus/example_dags/example_plexus.py → ...system/providers/plexus/example_plexus.py
@@ -15,33 +15,41 @@
 # specific language governing permissions and limitations
 # under the License.
 
+import os
 from datetime import datetime
 
 from airflow import DAG
 from airflow.providers.plexus.operators.job import PlexusJobOperator
 
 HOME = '/home/acc'
 T3_PRERUN_SCRIPT = 'cp {home}/imdb/run_scripts/mlflow.sh {home}/ && chmod +x mlflow.sh'.format(home=HOME)
+ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID")
+DAG_ID = "test"
 
-
-dag = DAG(
-    'test',
+with DAG(
+    DAG_ID,
     default_args={'owner': 'core scientific', 'retries': 1},
     description='testing plexus operator',
     start_date=datetime(2021, 1, 1),
     schedule_interval='@once',
     catchup=False,
-)
+) as dag:
+    # [START plexus_job_op]
+    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,
+        },
+    )
+    # [END plexus_job_op]
+
+
+from tests.system.utils import get_test_run  # noqa: E402
 
-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,
-)
+# Needed to run the example DAG with pytest (see: tests/system/README.md#run_via_pytest)
+test_run = get_test_run(dag)