playbook.yaml

- name: "Install build dependencies"
  hosts: localhost
  connection: local

  # The playbook requires passing 3 variables explicitly:
  # - stage: build or release. Different packages are installed depending on
  #          the chosen stage.
  # - arch: aarch64 or amd64. Architecture of the built image and wheels.
  # - accelerator: tpu or cuda. Available accelerator.
  pre_tasks:
    - name: "Validate required variables"
      ansible.builtin.assert:
        that: "{{ lookup('ansible.builtin.vars', item.name) is regex(item.pattern) }}"
        fail_msg: |
          "Variable '{{ item.name }} = '{{ lookup('ansible.builtin.vars', item.name) }}' doesn't match pattern '{{ item.pattern }}'"
          "Pass the required variable with: --e \"{{ item.name }}=<value>\""
      loop:
        - name: stage
          pattern: ^(build|release)$
        - name: arch
          pattern: ^(aarch64|amd64)$
        - name: accelerator
          pattern: ^(tpu|cuda)$

    - name: "Include vars from config files"
      ansible.builtin.include_vars:
        file: "config/{{ item }}"
      loop:
        # vars.yaml should be the first as other config files depend on it.
        - vars.yaml
        # cuda_deps should be loaded before apt, since apt depends on it.
        - cuda_deps.yaml
        - apt.yaml
        - pip.yaml
        - env.yaml
      tags: always # Execute this task even with `--skip-tags` or `--tags` is used.

  roles:
    - role: bazel
      tags: bazel

    - role: install_deps
      vars:
        apt_keys: "{{ apt.signing_keys }}"

        # If a variable (like `apt.pkgs.common`) is defined, but not set to
        # anything it cannot be concatenated with a list.
        # Use `v | default([], true)` to set `v` to an empty array if it evaluates to false.
        # See https://jinja.palletsprojects.com/en/3.0.x/templates/#jinja-filters.default.
        apt_pkgs: "{{
            apt.pkgs[stage + '_common'] | default([], true) +
              apt.pkgs[stage + '_' + arch] | default([], true) +
              apt.pkgs[stage + '_' + accelerator] | default([], true)
          }}"

        apt_repos: "{{ apt.repos }}"

        pip_pkgs: "{{
            pip.pkgs[stage + '_common'] | default([], true) +
            pip.pkgs[stage + '_' + arch] | default([], true) +
            pip.pkgs[stage + '_' + accelerator] | default([], true)
          }}"

        pip_pkgs_nodeps: "{{
            pip.pkgs_nodeps[stage + '_common'] | default([], true) +
            pip.pkgs_nodeps[stage + '_' + arch] | default([], true) +
            pip.pkgs_nodeps[stage + '_' + accelerator] | default([], true)
          }}"
      tags: install_deps

    - role: fetch_srcs
      vars:
        src_root: "/src"
      tags: fetch_srcs

    - role: build_srcs
      vars:
        src_root: "/src"
        env_vars: "{{
            build_env.common | default({}, true) |
            combine(build_env[arch] | default({}, true)) |
            combine(build_env[accelerator] | default({}, true))
          }}"
      tags: build_srcs

    - role: configure_env
      vars:
        env_vars: "{{
            release_env.common | default({}, true) |
            combine(release_env[arch] | default({}, true)) |
            combine(release_env[accelerator] | default({}, true))
          }}"
      when: stage == "release"
      tags: configure_env