Lautaro Celli

Posted on May 29

Implementando Databricks Asset Bundles sin morir en el intento

#databricks #dataengineering #cicd #azure

Índice

Introducción
¿Qué es Databricks Asset Bundles?
Antes de DABs, ¿es tan caótico?
Durante: implementando DABs en un proyecto productivo
Después de DABs: así se ve el paraíso
¿Tips & tricks?
Conclusión

Aclaración: esta es mi primera experiencia trabajando en Databricks, puede que más de una cosa pueda hacerse mejor y existan mejores prácticas, lo que voy a escribir es una experiencia propia.

Introducción

Hace unos meses empecé a trabajar con Databricks, en un proyecto "chico" para una empresa grande (consultoras ftw ¿?). Este proyectito era una pavadita: consumir data de archivos excel (para sorpresa de nadie), pegarle a una API, un par de transformaciones y volcar los datos en una Delta table. Bastante sencillo, así que resolvimos con un par de notebooks orquestadas por un Workflow en Databricks y salió andando. Hasta ahí, sencillísimo. El "problema" (o el desafío) vino cuando tuvimos que empezar a customizar bastante el pipeline, de originalmente 2 notebooks, para que usuarios de otros países puedan pasar a utilizar la herramienta que habíamos desarrollado. El país con el que trabajamos pasó de 1 a 6, se sumaron validaciones, nuevas reglas de negocio, diferencias entre ambientes y APIs para cada país, features nuevas como notificaciones custom directamente al usuario, export de resultados, nuevos flujos de datos, y un largo etcétera.

A esta altura, el workflow ya no era un único con 2 notebooks que funcionaban para un solo país. Ahora teníamos algo un poco más complejo, con parametrizaciones de por medio, validaciones de datos, notificaciones según estado y también nuevos workflows. Más allá de que los requerimientos estaban resueltos y teníamos algo funcional, ahora vinieron nuevos problemas: varias personas del equipo trabajando sobre los workflows, realizando modificaciones y teniendo poca o nula trazabilidad de los cambios en los pipelines.

Por suerte, la tecnología generalmente suele ofrecernos soluciones a este tipo de problemas comunes y así, tarde aunque aún a tiempo, implementamos Databricks Asset Bundles en nuestro proyecto.

¿Qué es Databricks Asset Bundles?

Databricks Asset Bundles (DABs) es una herramienta cuyo propósito es facilitar la adopción de las mejores prácticas de la ingeniería de software (versionado de código, testing, CI/CD, despliegue en distintos ambientes, entre otras) a p*royectos de ciencia de datos e IA, facilitando el trabajo colaborativo en proyectos en desarrollo. Según **Databricks*: "A bundle is an end-to-end definition of a project, including how the project should be structured, tested, and deployed". A grandes rasgos, lo que permite DABs es definir cualquier recurso de Databricks en archivos .yaml en nuestro repositorio, bajo el enfoque infrastructure-as-code (IaC).

Antes de DABs, ¿es tan caótico?

Si bien Asset Bundles, según la documentación, está pensado para proyectos complejos (vaya a saber uno la definición de complejo en este contexto), en nuestro caso sin ser necesariamente un proyecto grande se hizo necesario una vez que el número de workflows y personas trabajando sobre estos dejó de ser 1.

En nuestro caso, el principal problema de no tener implementado Asset Bundles estaba en no poder manejar el versionado de los workflows. Aunque la UI es bastante noble para desarrollar nuevos workflows, no existe la posibilidad de ver qué cambios se fueron haciendo ni se tiene registro de eso, aumentando la posibilidad de discrepancias entre pipelines, no pudiendo hacer code-reviews, sin poder tener un pantallazo general del código teniendo que ver los detalles tarea por tarea, y otros demases.

Puede parecer que no, pero el sólo hecho de no tener versionado para los workflows fue motivo suficiente para avanzar con DABs, a pesar de que no tenía mucha experiencia en Databricks ni tampoco en IaC ni CI/CD (mucha prueba y error).

Flujo de trabajo antes de DABs (principalmente en workflows):

Desarrollo de código, trackeado en un repositorio
Modificaciones en workflows en la UI, aplicando cambios 1 a 1 utilizando la interfaz. Si tenes workflows similares tenes que ir uno por uno en la UI, modificando 1 a 1, incluso tarea por tarea.

Durante: implementando DABs en un proyecto productivo

En nuestro caso, no necesitamos la parte de despliegue de clusters ni hacer muchas cosas raras, sólo precisamos -nuevamente- versionado de workflows, sumado a la validación y despliegue continuos de estos.

Debo decir que la documentación oficial de Databricks sobre cómo implementar DABs es bastante clara, aunque hubieron varias cuestiones no mencionadas explícitamente con las que me estuve peleando algunas horas hasta que tuve éxito.

Primero lo primero

Debemos tener instalada y configurada la CLI de Databricks. Se puede chequear fácilmente con

databricks --version

Una vez listo esto, es tan fácil como hacer

databricks bundle init

Y seguir los pasos del interactivo, que no es más que nombrar el bundle y elegir el template a utilizar. La estructura del bundle será la siguiente:

README.md  databricks.yml  fixtures  pytest.ini  requirements-dev.txt  resources  scratch  setup.py  src  tests

Mientras que databricks.yml se ve algo así

bundle:
  name: dabs101

include:
  - resources/*.yaml

targets:
  dev:
    # The default target uses 'mode: development' to create a development copy.
    # - Deployed resources get prefixed with '[dev my_user_name]'
    # - Any job schedules and triggers are paused by default.
    # See also https://docs.databricks.com/dev-tools/bundles/deployment-modes.html.
    mode: development
    default: true
    workspace:
      host: https://<development-workspace-url>

  prod:
    mode: production
    workspace:
      host: https://<production-workspace-url>
    permissions:
      - user_name: user1@mail.com
        level: CAN_MANAGE
    run_as:
      user_name: user2@mail.com

En este acotado yaml hay un par de variables a resaltar:

target: diferentes ambientes. Podemos definir más de un ambiente y, a su vez, utilizar distintos parámetros. Esto es útil para hacer despliegues en ambientes de dev o QA y hacer pruebas antes de un despliegue en prod (el que esté libre de pecado que tire la primera piedra).
include: ruta de los recursos incluidos en el bundle.

Para ver todas las posibilidades de configuración ver Databricks Asset Bundle configuration

Ahora si, vamos a lo importante: ¿cómo se vería la estructura de un workflow en yaml?

resources:
  jobs:
    foo_job:
      name: foo_job
      max_concurrent_runs: 1
      schedule:
        quartz_cron_expression: 12 0 0 * * ?
        timezone_id: America/Mexico_City
        pause_status: UNPAUSED
      tasks:
        - task_key: foo_task
          notebook_task:
            notebook_path: the/notebook/path
            source: GIT
          existing_cluster_id: existing_cluster_id
      git_source:
        git_url: git_url
        git_provider: git_provider
        git_branch: git_branch
      queue:
        enabled: true
      parameters:
        - name: parameter1
          default: value1
        - name: parameter2
          default: value2
        - name: parameter3
          default: ""

Pequeño truquito: si bien el yaml es sencillo de leer y saber qué hace, para desarrollar a mi me da bastante fiaca. Por suerte, se pueden desarrollar los workflows a través de la UI de Databricks y de ahí extraer el código fuente en yaml, json o código python.

Para la parte interesante que vendría siendo el despliegue del bundle, la CLI de Databricks dispone de 3 comandos bastante piolas

databricks bundle validate --target prod  # para validar el bundle e  identificar errores y warnings
databricks bundle summary --target prod   # resumen de lo que se va a desplegar (nuevo y/o asociado a recursos existentes)
databricks bundle deploy --target prod    # deploy del bundle

Hasta acá todo muy bueno para desplegar nuevos recursos, pero en nuestro caso los workflows ya existían pues habían sido desarrollados utilizando la UI de Databricks y al hacer el deploy los workflows se volvían a crear, teniendo jobs duplicados.

Me pelee bastante con comandos, archivos de metadata, entradas en los yamls y varias horas de trabajo, y al final la magia la hacía el comando bind

databricks bundle deployment unbind --target prod workflow_unique_name   # necesario si el recurso ya está asociado a un nuevo id en el bundle
databricks bundle deployment bind --target prod workflow_unique_name workflow_id

Con bind lo que hacemos es decirle a Databricks que el recurso existente de id X ahora pasa a ser gestionado por Asset Bundles: magia. Un detalle, importante prestar atención a la flag --target, ya que si no especifican lo hace en la por defecto y les puede pasar como a mi que no entienden por qué se sigue duplicando el workflow en prod cuando lo bindearon al entorno de dev.

Después de DABs: así se ve el paraíso

Bueno, tampoco para tanto. Como dije antes, nuestro principal problema ya no es tal, ahora tenemos versionado de los workflows.

Cabe destacar que mi motor fueron pura y exclusivamente los workflows, pero podemos implementar DABs para gestionar muchos más recursos en Databricks: worfklows, pipelines de Delta Live Tables, creación y despliegue de clusters, etc. Yo aproveché la inquietud para poder testear la herramienta por primera vez, aunque seguramente a partir de acá siga profundizando.

¿Tips & tricks?

En mis intentos por implementar DABs me crucé con varias cuestiones, un resumen para que dentro de lo posible se ahorren un poco más de tiempo que yo:

databricks bundle generate: en la teoría, este comando genera el .yaml de un recurso ya existente. En mi caso, los workflows estaban apuntados a notebooks conectados a un repositorio de Azure DevOps, por lo que las rutas eran relativas al repositorio. Por este motivo, el comando arrojaba error de path o formato de path incorrecto. Mi solución fue como comenté más arriba, directamente crear los archivos .yaml exportando el código de cada workflow, un poco más manual pero igual de válido.
Jobs duplicados: probé de diferentes maneras asociar los recursos existentes mediante nombres, id, rutas, etc., y si se te duplican los jobs es porque no están asociados correctamente al bundle. Hay 2 opciones: mantener el nuevo job gestionado por bundles y borrar el anterior (se pierde el run history) o bindear manualmente los recursos al bundle databricks bundle deployment bind/unbind: importante prestar atención a qué target estamos apuntando. Me pasó de tener configurado por defecto el target dev, por lo que al hacer bind seguía duplicando jobs, atención al flag --target. Algo importante es que si ya desplegaron el bundle y no están correctamente asociados los recursos, van a tener que hacer el unbind de cada recurso.
Aprovechar los pipelines de CI/CD: el despliegue de los bundles es idempotente, cada vez que se hace deploy se comparan los hashes entre los recursos desplegados y los del bundle, en caso de que no haya modificaciones no vuelve a desplegar el recurso, por lo que se puede dejar previsto la validación y despliegue del bundle en nuestro pipeline de CI/CD del repo sin preocuparnos por estar sobrescribiendo nada innecesariamente.
Una vez que los workflows pasan a ser gestionados por bundles pasa a bloquearse la edición en la UI, reforzando el concepto de IaC y evitar sobrescribirlos manualmente. Esto puede ser bastante fiaca si estamos en una etapa de constantes modificaciones donde sería más sencillo iterar trabajando en la UI que sobre los yaml, aunque siempre está disponible la opción de quitar el recurso del bundle para poder editar sobre la UI y luego volver a agregarlo al bundle (aunque no tendría tanto sentido pues iría contra el por qué de implementar DABs, pero la opción está).

Conclusión

Como cierre unas últimas preguntas y respuestas que me surgieron en este proceso:

¿Implementaría DABs en otros proyectos? Definitivamente si. Poder gestionar el despliegue completo desde un único archivo (o un conjunto de archivos yaml) me parece un puntazo a favor, además de poder tener todo esto versionado y trackeado en un repositorio
¿Lo implementaría en un proyecto desde 0? Si y no. Lo más fiaca es lo que comentaba más arriba sobre el blocker que significa en la UI para los workflows. Creo que se podría implementar desde 0 en el proyecto, pero para recursos que necesitemos iterar bastante como puede ser un workflow o un pipeline en escenarios tempranos de un desarrollo podrían sumarse al último o cuando tengan un grado de madurez suficiente como para no necesitar estar aplicando cambios constantemente en archivos yaml.
La comunicación del bundle no es bidireccional: si modificamos algo en la UI no se modifica automáticamente en el bundle. Es importante tener esto en cuenta pues si modificamos recursos a través de la UI y no reflejamos los cambios en la definición de nuestro bundle, cada vez que hagamos un deploy vamos a pisar la configuración según lo que tengamos definido. Debería poder implementarse algún pipeline para verificar estos cambios de estado en los recursos, aunque desconozco.

Para cerrar, nuevamente aclarar que esta fue la primera vez que trabajo con Databricks/Databricks Asset Bundles, por lo que muchas cosas seguro puedan hacerse de mejor manera. La idea no es que se tome esto como una guía per sé, si no que mi experiencia sirva de referencia para personas que por ahí estén en una situación parecida a la que estaba yo hace unos días sin tener mínima idea de cómo implementar la herramienta.