slurm.jinja

# Copyright 2017 SchedMD LLC.
# Modified for use with the Slurm Resource Manager.
#
# Copyright 2015 Google Inc. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

{% set _ = properties.update({
    'project': env['project'],
    'compute_node_prefix': properties["cluster_name"] ~ "-compute",
    }) %}
{% set cluster_region =  '-'.join(properties['zone'].split('-')[:-1]) %}
{% set def_network_name = properties["cluster_name"] ~ "-network" %}

{% set regions = {
          cluster_region: 
            (properties['vpc_subnet'] or (properties['cluster_name']~'-'~cluster_region),
             properties['vpc_subnet'])
        } %}
{% for part in properties['partitions'] %}
{% set region = '-'.join(part['zone'].split('-')[:-1]) %}
{% set _ = regions.update({region:
            (part['vpc_subnet'] or properties['cluster_name']~'-'~region,
             regions.get(region, None)[1] or part['vpc_subnet'])
          }) %}
{% set _ = part.update({'region': region}) %}
{% endfor %}

resources:
{% if not properties['shared_vpc_host_project'] %}
{% if not properties['vpc_net'] %}
- name: {{def_network_name}}
  type: compute.v1.network
  properties:
    autoCreateSubnetworks: false
{% endif %}

{% for region, (subnet, existing) in regions.items() if not existing %}
- name: {{subnet}}
  type: compute.v1.subnetwork
  properties:
  {% if properties['vpc_net'] %}
    network: https://www.googleapis.com/compute/v1/projects/{{properties['project']}}/global/networks/{{properties['vpc_net']}}
  {% else %}
    network: $(ref.{{def_network_name}}.selfLink)
  {% endif %}
    ipCidrRange: {{ "10.{}.0.0/16".format(loop.index0) }}
    region: {{region}}
    privateIpGoogleAccess: {{ properties["private_google_access"]}}
{% endfor %}

{% if not properties['vpc_net'] %}
{% if (properties['external_compute_ips']
       or properties['external_login_ips']
       or properties['external_controller_ip']) %}
- name: {{properties["cluster_name"]}}-allow-ssh
  type: compute.v1.firewall
  properties:
    network: $(ref.{{def_network_name}}.selfLink)
    sourceRanges: ["0.0.0.0/0"]
    allowed:
    - IPProtocol: TCP
      ports: ["22"]
    - IPProtocol: ICMP
{% else %}
- name: {{properties["cluster_name"]}}-allow-iap
  type: compute.v1.firewall
  properties:
    network: $(ref.{{def_network_name}}.selfLink)
    sourceRanges: ["35.235.240.0/20"]
    allowed:
    - IPProtocol: TCP
      ports: ["22"]
    - IPProtocol: ICMP
{% endif %}
- name: {{properties["cluster_name"]}}-all-internal-firewall-rule
  type: compute.v1.firewall
  properties:
    network: $(ref.{{def_network_name}}.selfLink)
    sourceRanges: ["10.0.0.0/8"]
    allowed:
    - IPProtocol: TCP
      ports: ["0-65535"]
    - IPProtocol: UDP
      ports: ["0-65535"]
    - IPProtocol: ICMP
{% endif %}

{% if (not properties['shared_vpc_host_project']) and
       not (properties['external_compute_ips'] and
            properties['external_login_ips'] and
            properties['external_controller_ip']) %}
{% for region, (subnet, existing) in regions.items() %}
- name: {{properties["cluster_name"]}}-{{region}}-router
  type: compute.v1.router
  properties:
  {% if properties['vpc_net'] %}
    network: https://www.googleapis.com/compute/v1/projects/{{ properties["project"] }}/global/networks/{{ properties["vpc_net"]}}
  {% else %}
    network: $(ref.{{def_network_name}}.selfLink)
  {% endif %}
    region: {{region}}
    nats:
    - name: {{properties["cluster_name"]}}-nat
      natIpAllocateOption: "AUTO_ONLY"
      sourceSubnetworkIpRangesToNat: "LIST_OF_SUBNETWORKS"
      subnetworks:
    {% if existing %}
      - name: https://www.googleapis.com/compute/v1/projects/{{properties['project']}}/regions/{{region}}/subnetworks/{{subnet}}
    {% else %}
      - name: $(ref.{{subnet}}.selfLink)
    {% endif %}
        sourceIpRangesToNat: ["PRIMARY_IP_RANGE"]
{% endfor %}
{% endif %}
{% endif %}

- name: {{properties["cluster_name"]}}-controller
  type: compute.v1.instance
  properties:
    zone: {{ properties["zone"] }}
    machineType: https://www.googleapis.com/compute/v1/projects/{{ properties["project"] }}/zones/{{ properties["zone"] }}/machineTypes/{{ properties["controller_machine_type"] }}
    disks:
    - deviceName: boot
      type: PERSISTENT
      boot: true
      autoDelete: true
      initializeParams:
        sourceImage: https://www.googleapis.com/compute/v1/projects/centos-cloud/global/images/family/centos-7
        diskType: https://www.googleapis.com/compute/v1/projects/{{ properties["project"] }}/zones/{{ properties["zone"] }}/diskTypes/{{ properties["controller_disk_type"] }}
        diskSizeGb: {{ properties["controller_disk_size_gb"] }}
{% if properties['controller_secondary_disk'] %}
    - deviceName: secondary
      type: PERSISTENT
      autoDelete: TRUE
      initializeParams:
        diskType: https://www.googleapis.com/compute/v1/projects/{{ properties["project"] }}/zones/{{ properties["zone"] }}/diskTypes/{{ properties["controller_secondary_disk_type"] }}
        diskSizeGb: {{properties["controller_secondary_disk_size_gb"]}}
{% endif %}
{% if properties ['shared_vpc_host_project'] %}
    networkInterfaces:
    - subnetwork: https://www.googleapis.com/compute/v1/projects/{{ properties["shared_vpc_host_project"] }}/regions/{{cluster_region}}/subnetworks/{{ properties['vpc_subnet'] }}
{% elif properties['vpc_subnet'] %}
    networkInterfaces:
    - subnetwork: https://www.googleapis.com/compute/v1/projects/{{ properties["project"] }}/regions/{{cluster_region}}/subnetworks/{{ properties['vpc_subnet'] }}
{% else %}
    networkInterfaces:
    - subnetwork: $(ref.{{regions[cluster_region][0]}}.selfLink)
{% endif %}
{% if properties['external_controller_ip'] %}
      accessConfigs:
      - name: External NAT
        type: ONE_TO_ONE_NAT
{% endif %}
    serviceAccounts:
      - email: {{ properties["controller_service_account"] }}
        scopes: {{ properties["controller_scopes"] }}
    tags:
      items:
        - controller
{% if properties["controller_labels"] %}
    labels:
      {{properties["controller_labels"]}}
{% endif %}
    metadata:
      items:
        - key: startup-script
          value: |
            {{ imports["scripts/startup.sh"]|indent(12) }}
        - key: util_script
          value: |
            {{ imports["scripts/util.py"]|indent(12) }}
        - key: config
          value: |
            {{ properties|safe|tojson(2)|indent(12) }}
        - key: setup_script
          value: |
            {{ imports["scripts/setup.py"]|indent(12) }}
        - key: slurm_resume
          value: |
            {{ imports["scripts/resume.py"]|indent(12) }}
        - key: slurm_suspend
          value: |
            {{ imports["scripts/suspend.py"]|indent(12) }}
        - key: slurmsync
          value: |
            {{ imports["scripts/slurmsync.py"]|indent(12) }}
        - key: enable-oslogin
          value: "TRUE"
        - key: VmDnsSetting
          value: GlobalOnly
        - key: custom-compute-install
          value: |
            {{ imports["scripts/custom-compute-install"]|indent(12) }}
        - key: custom-controller-install
          value: |
            {{ imports["scripts/custom-controller-install"]|indent(12) }}
        - key: compute-shutdown
          value: |
            {{ imports["scripts/compute-shutdown"]|indent(12) }}
        - key: slurm_conf_tpl
          value: |
            {{ imports["etc/slurm.conf.tpl"]|indent(12) }}
        - key: slurmdbd_conf_tpl
          value: |
            {{ imports["etc/slurmdbd.conf.tpl"]|indent(12) }}
        - key: cgroup_conf_tpl
          value: |
            {{ imports["etc/cgroup.conf.tpl"]|indent(12) }}
        - key: fluentd_conf_tpl
          value: |
            {{ imports["etc/controller-fluentd.conf.tpl"]|indent(12) }}


{% for n in range(properties['login_node_count']) %}
- name: {{properties["cluster_name"]}}-login{{ n }}
  type: compute.v1.instance
  properties:
    zone: {{ properties["zone"] }}
    machineType: https://www.googleapis.com/compute/v1/projects/{{ properties["project"] }}/zones/{{ properties["zone"] }}/machineTypes/{{ properties["login_machine_type"] }}
    disks:
    - deviceName: boot
      type: PERSISTENT
      boot: true
      autoDelete: true
      initializeParams:
        sourceImage: https://www.googleapis.com/compute/v1/projects/centos-cloud/global/images/family/centos-7
        diskType: https://www.googleapis.com/compute/v1/projects/{{ properties["project"] }}/zones/{{ properties["zone"] }}/diskTypes/{{ properties["login_disk_type"] }}
        diskSizeGb: {{ properties["login_disk_size_gb"] }}
  {% if properties ['shared_vpc_host_project'] %}
    networkInterfaces:
    - subnetwork: https://www.googleapis.com/compute/v1/projects/{{ properties["shared_vpc_host_project"] }}/regions/{{cluster_region}}/subnetworks/{{ properties['vpc_subnet'] }}
  {% elif properties['vpc_subnet'] %}
    networkInterfaces:
    - subnetwork: https://www.googleapis.com/compute/v1/projects/{{ properties["project"] }}/regions/{{cluster_region}}/subnetworks/{{ properties['vpc_subnet'] }}
  {% else %}
    networkInterfaces:
    - subnetwork: $(ref.{{regions[cluster_region][0]}}.selfLink)
  {% endif %}
{% if properties['external_login_ips'] %}
      accessConfigs:
      - name: External NAT
        type: ONE_TO_ONE_NAT
{% endif %}
    serviceAccounts:
      - email: {{ properties["login_node_service_account"] }}
        scopes:
          {{properties["login_node_scopes"]}}
    tags:
      items:
        - login
  {% if properties["login_labels"] %}
    labels:
      {{properties["login_labels"]}}
  {% endif %}
    metadata:
      items:
        - key: startup-script
          value: |
            {{ imports["scripts/startup.sh"]|indent(12) }}
        - key: util_script
          value: |
            {{ imports["scripts/util.py"]|indent(12) }}
        - key: config
          value: |
            {{ properties|safe|tojson(2)|indent(12) }}
        - key: setup_script
          value: |
            {{ imports["scripts/setup.py"]|indent(12) }}
        - key: enable-oslogin
          value: "TRUE"
        - key: VmDnsSetting
          value: GlobalOnly
{% endfor %}

{% for i in range(properties["partitions"]|length) %}
{% set region = properties['partitions'][i]['region'] %}

{% for n in range(properties["partitions"][i]['static_node_count']) %}
- name: {{properties["compute_node_prefix"]}}-{{i}}-{{n}}
  type: compute.v1.instance
  properties:
    zone: {{ properties["partitions"][i]["zone"] }}
    machineType: https://www.googleapis.com/compute/v1/projects/{{ properties["project"] }}/zones/{{ properties["partitions"][i]["zone"] }}/machineTypes/{{ properties["partitions"][i]["machine_type"] }}
  {% if properties["partitions"][i]["cpu_platform"] %}
    minCpuPlatform: {{ properties["partitions"][i]["cpu_platform"]}}
  {% endif %}
  {% if properties["partitions"][i]["gpu_type"] %}
    guestAccelerators:
    - acceleratorType: https://www.googleapis.com/compute/v1/projects/{{ properties["project"] }}/zones/{{ properties["partitions"][i]["zone"] }}/acceleratorTypes/{{ properties["partitions"][i]["gpu_type"] }}
      acceleratorCount: {{ properties["partitions"][i]['gpu_count']}}
    scheduling:
      onHostMaintenance: TERMINATE
  {% endif %}
    disks:
    - deviceName: boot
      type: PERSISTENT
      boot: true
      autoDelete: true
      initializeParams:
        sourceImage: https://www.googleapis.com/compute/v1/projects/centos-cloud/global/images/family/centos-7
        diskType: https://www.googleapis.com/compute/v1/projects/{{ properties["project"] }}/zones/{{ properties["partitions"][i]["zone"] }}/diskTypes/{{ properties["partitions"][i]["compute_disk_type"] }}
        diskSizeGb: {{ properties["partitions"][i]["compute_disk_size_gb"] }}
  {% if properties['shared_vpc_host_project']  %}
    networkInterfaces:
    - subnetwork: https://www.googleapis.com/compute/v1/projects/{{ properties["shared_vpc_host_project"] }}/regions/{{region}}/subnetworks/{{ properties['partitions'][i]['vpc_subnet'] }}
  {% elif properties['partitions'][i]['vpc_subnet'] %}
    networkInterfaces:
    - subnetwork: https://www.googleapis.com/compute/v1/projects/{{ properties["project"] }}/regions/{{region}}/subnetworks/{{ properties['partitions'][i]['vpc_subnet'] }}
  {% else %}
    networkInterfaces:
    - subnetwork: $(ref.{{regions[region][0]}}.selfLink)
  {% endif %}
  {% if properties['external_compute_ips'] %}
      accessConfigs:
      - name: External NAT
        type: ONE_TO_ONE_NAT
  {% endif %}
    serviceAccounts:
      - email: {{ properties["compute_node_service_account"] }}
        scopes:
          {{properties["compute_node_scopes"]}}
    tags:
      items:
        - compute
  {% if properties["partitions"][i]["compute_labels"] %}
    labels:
      {{properties["partitions"][i]["compute_labels"]}}
  {% endif %}
    metadata:
      items:
        - key: startup-script
          value: |
            {{ imports["scripts/startup.sh"]|indent(12) }}
        - key: util_script
          value: |
            {{ imports["scripts/util.py"]|indent(12) }}
        - key: config
          value: |
            {{ properties|safe|tojson(2)|indent(12) }}
        - key: setup_script
          value: |
            {{ imports["scripts/setup.py"]|indent(12) }}
        - key: enable-oslogin
          value: "TRUE"
        - key: VmDnsSetting
          value: GlobalOnly
        - key: fluentd_conf_tpl
          value: |
            {{ imports["etc/compute-fluentd.conf.tpl"]|indent(12) }}
{% endfor %}
{% endfor %}

{% for i in range(properties["partitions"]|length) %}
- name: {{properties["compute_node_prefix"]}}-{{i}}-image
  type: compute.v1.instance
  properties:
    zone: {{ properties["zone"] }}
    machineType: https://www.googleapis.com/compute/v1/projects/{{ properties["project"] }}/zones/{{ properties["zone"] }}/machineTypes/{{ properties["compute_image_machine_type"] }}
    disks:
    - deviceName: boot
      type: PERSISTENT
      boot: true
      autoDelete: true
      initializeParams:
        sourceImage: https://www.googleapis.com/compute/v1/projects/centos-cloud/global/images/family/centos-7
        diskType: https://www.googleapis.com/compute/v1/projects/{{ properties["project"] }}/zones/{{ properties["zone"] }}/diskTypes/{{ properties["compute_image_disk_type"] }}
        diskSizeGb: {{ properties["compute_image_disk_size_gb"] }}
  {% if properties ['shared_vpc_host_project'] %}
    networkInterfaces:
    - subnetwork: https://www.googleapis.com/compute/v1/projects/{{ properties["shared_vpc_host_project"] }}/regions/{{cluster_region}}/subnetworks/{{ properties['vpc_subnet'] }}
  {% elif properties['vpc_subnet'] %}
    networkInterfaces:
    - subnetwork: https://www.googleapis.com/compute/v1/projects/{{ properties["project"] }}/regions/{{cluster_region}}/subnetworks/{{ properties['vpc_subnet'] }}
  {% else %}
    networkInterfaces:
    - subnetwork: $(ref.{{regions[cluster_region][0]}}.selfLink)
  {% endif %}
  {% if properties['external_compute_ips'] %}
      accessConfigs:
      - name: External NAT
        type: ONE_TO_ONE_NAT
  {% endif %}
    serviceAccounts:
      - email: "default"
        scopes:
        - "https://www.googleapis.com/auth/cloud-platform"
    tags:
      items:
        - compute
  {% if properties["compute_image_labels"] %}
    labels:
      {{properties["compute_image_labels"]}}
  {% endif %}
    metadata:
      items:
        - key: startup-script
          value: |
            {{ imports["scripts/startup.sh"]|indent(12) }}
        - key: util_script
          value: |
            {{ imports["scripts/util.py"]|indent(12) }}
        - key: config
          value: |
            {{ properties|safe|tojson(2)|indent(12) }}
        - key: setup_script
          value: |
            {{ imports["scripts/setup.py"]|indent(12) }}
        - key: enable-oslogin
          value: "TRUE"
        - key: VmDnsSetting
          value: GlobalOnly
        - key: fluentd_conf_tpl
          value: |
            {{ imports["etc/compute-fluentd.conf.tpl"]|indent(12) }}
{% endfor %}