[AGENT-5360] Create Airflow Operator for Custom Job

datarobot_provider/__init__.py

@@ -7,10 +7,10 @@
 # Released under the terms of DataRobot Tool and Utility Agreement.
 def get_provider_info():
     return {
-        "package-name": "airflow-provider-datarobot",
+        "package-name": "airflow-provider-datarobot-early-access",
         "name": "DataRobot Airflow Provider",
         "description": "DataRobot Airflow provider.",
-        "versions": ["0.0.8"],
+        "versions": ["0.0.8.1"],
         "connection-types": [
             {
                 "hook-class-name": "datarobot_provider.hooks.datarobot.DataRobotHook",

datarobot_provider/example_dags/custom_job/job.py

@@ -0,0 +1,31 @@
+import os
+from datarobot import Deployment
+
+
+def main():
+    print(f"Running python code: {__file__}")
+
+    # Using this job runtime parameters
+    print()
+    print("Runtime parameters:")
+    print("-------------------")
+    string_param = os.environ.get("STRING_PARAMETER", None)
+    print(f"string param: {string_param}")
+
+    deployment_param = os.environ.get("DEPLOYMENT", None)
+    print(f"deployment_param: {deployment_param}")
+
+    model_package_param = os.environ.get("MODEL_PACKAGE", None)
+    print(f"model_package_param: {model_package_param}")
+
+    # An example of using the python client to list deployments
+    deployments = Deployment.list()
+    print()
+    print("List of all deployments")
+    print("-----------------------")
+    for deployment in deployments:
+        print(deployment)
+
+
+if __name__ == "__main__":
+    main()

datarobot_provider/example_dags/custom_job/metadata.yaml

@@ -0,0 +1,12 @@
+name: runtime-params
+
+runtimeParameterDefinitions:
+  - fieldName: MODEL_PACKAGE
+    type: modelPackage
+    description: Model package that will be used to store key values
+  - fieldName: DEPLOYMENT
+    type: deployment
+    description: Deployment that will be used to make predictions
+  - fieldName: STRING_PARAMETER
+    type: string
+    description: An example of a string parameter

datarobot_provider/example_dags/custom_job/run.sh

@@ -0,0 +1,60 @@
+#!/bin/bash
+
+echo "Job Starting: ($0)"
+
+echo "===== Runtime Parameters ======"
+echo "Model Package:     $MODEL_PACKAGE"
+echo "Deployment:        $DEPLOYMENT"
+echo "STRING_PARAMETER:  $STRING_PARAMETER"
+echo
+echo
+echo "===== Generic Variables ==========================="
+echo "CURRENT_CUSTOM_JOB_RUN_ID: $CURRENT_CUSTOM_JOB_RUN_ID"
+echo "CURRENT_CUSTOM_JOB_ID:     $CURRENT_CUSTOM_JOB_ID"
+echo "DATAROBOT_ENDPOINT:        $DATAROBOT_ENDPOINT"
+echo "DATAROBOT_API_TOKEN:       Use the environment variable $DATAROBOT_API_TOKEN"
+echo "==================================================="
+
+echo
+echo "How to check how much memory your job has"
+  memory_limit_bytes=$(cat /sys/fs/cgroup/memory/memory.limit_in_bytes)
+  memory_limit_megabytes=$((memory_limit_bytes / 1024 / 1024))
+echo "Memory Limit (in Megabytes): $memory_limit_megabytes"
+echo
+
+# Uncomment the following if you want to check if the job has network access
+## Define the IP address of an external server to ping (e.g., Google's DNS)
+#external_server="8.8.8.8"
+#echo "Checking internet connection"
+## Try to ping the external server
+#ping -c 1 $external_server > /dev/null 2>&1
+#
+## Check the exit status of the ping command
+#if [ $? -eq 0 ]; then
+#    echo "Internet connection is available."
+#else
+#    echo "No internet connection."
+#fi
+#echo
+#echo
+
+# Run the code in job.py
+dir_path=$(dirname $0)
+echo "Entrypoint is at $dir_path - cd into it"
+cd $dir_path
+
+if command -v python3 &>/dev/null; then
+    echo "python3 is installed and available."
+else
+    echo "Error: python3 is not installed or not available."
+    exit 1
+fi
+
+python_file="job.py"
+if [ -f "$python_file" ]; then
+    echo "Found $python_file .. running it"
+    python3 ./job.py
+else
+    echo "File $python_file does not exist"
+    exit 1
+fi

datarobot_provider/example_dags/datarobot_custom_job_dag.py

@@ -0,0 +1,92 @@
+# Copyright 2024 DataRobot, Inc. and its affiliates.
+#
+# All rights reserved.
+#
+# This is proprietary source code of DataRobot, Inc. and its affiliates.
+#
+# Released under the terms of DataRobot Tool and Utility Agreement.
+
+from datetime import datetime
+
+from airflow.decorators import dag
+
+from datarobot_provider.operators.custom_job import CreateCustomJobOperator
+from datarobot_provider.operators.custom_job import AddFilesToCustomJobOperator
+from datarobot_provider.operators.custom_job import SetCustomJobExecutionEnvironmentOperator
+from datarobot_provider.operators.custom_job import SetCustomJobRuntimeParametersOperator
+from datarobot_provider.operators.custom_job import RunCustomJobOperator
+from datarobot_provider.sensors.client import BaseAsyncResolutionSensor
+
+
+@dag(
+    schedule=None,
+    start_date=datetime(2023, 1, 1),
+    tags=['example', 'custom job'],
+    params={},
+)
+def create_custom_custom_job():
+    create_custom_job_op = CreateCustomJobOperator(
+        task_id='create_custom_job',
+        name="airflow-test-create-custom-job-v556",
+        description="demo-test-demonstration",
+    )
+
+    add_files_to_custom_job_op = AddFilesToCustomJobOperator(
+        task_id='add_files_to_custom_job',
+        custom_job_id=create_custom_job_op.output,
+        files_path="custom_job/",
+    )
+
+    # list_execution_env_op = ListExecutionEnvironmentOperator(
+    #     task_id='list_execution_env',
+    #     search_for="Python 3.9 PyTorch Drop-In"
+    # )
+
+    set_env_to_custom_job_op = SetCustomJobExecutionEnvironmentOperator(
+        task_id='set_env_to_custom_job',
+        custom_job_id=create_custom_job_op.output,
+        environment_id='5e8c888007389fe0f466c72b',
+        environment_version_id='65c1db901800cd9782d7ac07',
+    )
+
+    set_runtime_parameters_op = SetCustomJobRuntimeParametersOperator(
+        task_id='set_runtime_parameters',
+        custom_job_id=create_custom_job_op.output,
+        runtime_parameter_values=[
+            {"fieldName": "DEPLOYMENT", "type": "deployment", "value": "650ef15944f21ea1a3c91a25"},
+            {
+                "fieldName": "MODEL_PACKAGE",
+                "type": "modelPackage",
+                "value": "654b9b228404a39b5c8da5b2",
+            },
+            {"fieldName": "STRING_PARAMETER", "type": "string", "value": 'my test string'},
+        ],
+    )
+
+    run_custom_job_op = RunCustomJobOperator(
+        task_id='run_custom_job',
+        custom_job_id=create_custom_job_op.output,
+    )
+
+    custom_job_complete_sensor = BaseAsyncResolutionSensor(
+        task_id="check_custom_job_complete",
+        job_id=run_custom_job_op.output,
+        poke_interval=5,
+        mode="reschedule",
+        timeout=3600,
+    )
+
+    (
+        create_custom_job_op
+        >> add_files_to_custom_job_op
+        >> set_env_to_custom_job_op
+        >> set_runtime_parameters_op
+        >> run_custom_job_op
+        >> custom_job_complete_sensor
+    )
+
+
+create_custom_job_dag = create_custom_custom_job()
+
+if __name__ == "__main__":
+    create_custom_job_dag.test()