Package: neuro-airflow-plugin
Neu.ro Airflow pluginThis package helps you execute your ML workloads onneu.roplatform from Airflow environment.Also, take a look at ourCLI referenceandPython API reference.EnvironmentPython 3.6+apache-airflow >= 1.10.xNeuromation >= 20.4.6InstallationThe plugin is written to automatically register with Airflow, so all you have to do is install it into your Python environment:pip install neuro-airflow-pluginUsageBefore start you need to get a Neuro token by using CLI command:neuro config show-tokenand set up a Neuro Connection (neuro_defaultby default) in Airflow:airflow connections --add \
    --conn_id neuro_default2 \
    --conn_type "neuro" \
    --conn_extra '{"token": "Put your Token here..."}'Apart fromtokenyou can also provide those fields as part of extra json:cluster- name of the cluster used for compute scheduling. Default cluster will be used if not provided.api_url- entry URL for Neuro Platform. Only needed for custom clusters.You can set up the connection from UI interface as well, just put the same JSON document
intoExtraform field. Connection type does not matter, so you can pick any that the UI allows.For more information on how to set up connections in Airflow seeManaging Connections.NeuroRunOperatorRuns a Job in the Neuro Platform. Example usage:fromairflow.operators.neuroimportNeuroRunOperatorrun=NeuroRunOperator(task_id="small-deeds",job_command="echo 'Big things start from small deeds'",job_image="ubuntu:{{ var.value.ubuntu_tag }}",job_resources_preset="{% if var.value.use_large_cpu %}cpu-large{% else %}cpu-small{% endif %}")For more usage examples seeexamples/dagsfolder of the repository.Operator argumentsjob_commandstrRequired- Command to be executed in the Job. If you need to override the
entrypoint of an image, seejob_entrypointinstead.job_imagestrRequired- Container image used for the Job. Name can be either a docker image
name hosted on an external public repository or a Neuro image specified byimage://scheme.job_namestr- Optional job name. Note that creating 2 running jobs with the same name by the
same user is forbidden.job_volumeslist- List of strings describing a volume mount orneuromation.Volumeobjects.
String description consists of 3 parts separated by column:storage URI,mount path,mount mode.
For example:storage:my-project:/var/storage:ro.job_resources_presetstr- Predefined resource configuration (to see available values, runneuro config show)job_resourcesResources- Custom resource configuration. SeePython API referencefor details.job_is_preemptiblebool- Whether the Job may be run on a preemptible, or also known as Spot instance.
Is only used with custom resource configuration.job_extshmbool- Request extended '/dev/shm' space. Defaults toTrueand is only used with
predefined resource configuration.job_tagslist- List of string tags to mark the Job with. Can later be used for filtering, etc.job_descriptionstr- Optional job description in free format.job_lifespanfloat- Optional job run-time limit in seconds. Is unlimited by default.job_environdict- Environment variables to run the Job with. Jinja template support is only provided for values,
not for keys, see more details below.job_entrypointstr- Override ENTRYPOINT of the container image.job_http_portstr- Enable HTTP port forwarding to specified container port. If used you can access it from
a custom link definition on the Task panel in Airflow UI (seeAirflow docsfor details on how it works). Disabled by default.job_http_authbool- Disable Neuro authentication on the exposed port injob_http_port.job_ttybool- Allocate a TTY for the Container.job_detachbool- Detach after starting the job. If detached Job logs will not be viewable in Airflow interface,
but the job will not consume Airflow worker slot. Defaults toTrue.raise_on_errnobool- Raise an error if job returns a non-zero exit code. Ignored ifjob_detachisTrue. Default toTrue.neuro_conn_idbool- Name of the connection to use for Neuro authentication. Defaults toneuro_default.See also theneuro runreference inCLI documentationJinja2 template fieldsAirflow supports passing custom attributes and dynamic definitions using Jinja templating fields. This operator
supports templating on the following fields:job_commandjob_imagejob_volumesjob_namejob_resources_presetjob_tagsjob_environjob_entrypointjob_descriptionjob_http_portneuro_conn_idXCom exportsThe operator exports 2 XCom values:return_value(default in Airflow for query) andassigned_job. Both are
JSON documents with the following fields:idstr- Job ID assigned by Neuro on start.exit_codeint- Command return code if the Job already finished.statusstr- One of job statuses:pending,running,succeeded,failedorunknown.http_urlstr- URL of the exposed HTTP port ifjob_http_portis used.NeuroJobStatusSensorWait for a Job to be completed or any other status transition to happen. Example usage:fromairflow.sensors.neuroimportNeuroJobStatusSensorwait=NeuroJobStatusSensor(task_id="wait_close",job_id="{{ task_instance.xcom_pull(task_ids='small-deeds')['id'] }}",# noqapoke_interval=5,timeout=10*60,)Operator argumentsjob_idstr- Job ID to query for status updates.job_statuseslist- ListJobStatusenum values to wait for.neuro_conn_idstr- Name of the connection to use for Neuro authentication. Defaults toneuro_default.Jinja2 template fieldsjob_idXCom exportsDoes not export any XCom values.NeuroHookIn some cases you may need to access other functionalities of the platform. This can be done using the NeuroHook.
For example:importyarlfromneuromation.apiimportResourceNotFoundfromairflow.hooks.neuroimportNeuroHookfromairflow.operators.python_operatorimportBranchPythonOperatordefcheck_model(templates_dict,**kw):hook=NeuroHook()withhook:try:hook.run(hook.client.storage.stat(yarl.URL("storage:"+templates_dict["model_path"])))return"process_with_model"exceptResourceNotFound:return"process_without_model"check_model=BranchPythonOperator(task_id="check_model_exists",python_callable=check_model,provide_context=True,templates_dict={"model_path":"{{ var.value.project_home }}/model.pth"},)Explore thePython SDKfor more features of the Platform.