<commit_before>import os
import logging
import base64
import json

from hailtop import aiogoogle

from ..batch_configuration import PROJECT, DEFAULT_NAMESPACE
from ..worker_config import WorkerConfig
from ..log_store import LogStore

log = logging.getLogger('create_instance')

BATCH_WORKER_IMAGE = os.environ['HAIL_BATCH_WORKER_IMAGE']

log.info(f'BATCH_WORKER_IMAGE {BATCH_WORKER_IMAGE}')


async def create_instance(app, zone, machine_name, machine_type, activation_token,
                          max_idle_time_msecs, worker_local_ssd_data_disk,
                          worker_pd_ssd_data_disk_size_gb, boot_disk_size_gb):
    log_store: LogStore = app['log_store']
    compute_client: aiogoogle.ComputeClient = app['compute_client']

    if worker_local_ssd_data_disk:
        worker_data_disk = {
            'type': 'SCRATCH',
            'autoDelete': True,
            'interface': 'NVME',
            'initializeParams': {
                'diskType': f'zones/{zone}/diskTypes/local-ssd'
            }}
        worker_data_disk_name = 'nvme0n1'
    else:
        worker_data_disk = {
            'autoDelete': True,
            'initializeParams': {
                'diskType': f'projects/{PROJECT}/zones/{zone}/diskTypes/pd-ssd',
                'diskSizeGb': str(worker_pd_ssd_data_disk_size_gb)
            }
        }
        worker_data_disk_name = 'sdb'

    config = {
        'name': machine_name,
        'machineType': f'projects/{PROJECT}/zones/{zone}/machineTypes/{machine_type}',
        'labels': {
            'role': 'batch2-agent',
            'namespace': DEFAULT_NAMESPACE
        },

        'disks': [{
            'boot': True,
            'autoDelete': True,
            'initializeParams': {
                'sourceImage': f'projects/{PROJECT}/global/images/batch-worker-12',
                'diskType': f'projects/{PROJECT}/zones/{zone}/diskTypes/pd-ssd',
                'diskSizeGb': str(boot_disk_size_gb)
            }
        }, worker_data_disk],

        'networkInterfaces': [{
            'network': 'global/networks/default',
            'networkTier': 'PREMIUM',
            'accessConfigs': [{
                'type': 'ONE_TO_ONE_NAT',
                'name': 'external-nat'
            }]
        }],

        'scheduling': {
            'automaticRestart': False,
            'onHostMaintenance': "TERMINATE",
            'preemptible': True
        },

        'serviceAccounts': [{
            'email': f'<EMAIL>',
            'scopes': [
                'https://www.googleapis.com/auth/cloud-platform'
            ]
        }],

        'metadata': {
            'items': [{
                'key': 'startup-script',
                'value': '''
#!/bin/bash
set -x

curl -s -H "Metadata-Flavor: Google" "http://metadata.google.internal/computeMetadata/v1/instance/attributes/run_script"  >./run.sh

nohup /bin/bash run.sh >run.log 2>&1 &
    '''
            }, {
                'key': 'run_script',
                'value': rf'''
#!/bin/bash
set -x

# only allow udp/53 (dns) to metadata server
# -I inserts at the head of the chain, so the ACCEPT rule runs first
iptables -I DOCKER-USER -d 169.254.169.254 -j DROP
iptables -I DOCKER-USER -d 169.254.169.254 -p udp -m udp --destination-port 53 -j ACCEPT

docker network create public --opt com.docker.network.bridge.name=public
docker network create private --opt com.docker.network.bridge.name=private
# make the internal network not route-able
iptables -I DOCKER-USER -i public -d 10.0.0.0/8 -j DROP
# make other docker containers not route-able
iptables -I DOCKER-USER -i public -d 172.16.0.0/12 -j DROP
# not used, but ban it anyway!
iptables -I DOCKER-USER -i public -d 192.168.0.0/16 -j DROP

WORKER_DATA_DISK_NAME="{worker_data_disk_name}"

# format local SSD
sudo mkfs.xfs -m reflink=1 /dev/$WORKER_DATA_DISK_NAME
sudo mkdir -p /mnt/disks/$WORKER_DATA_DISK_NAME
sudo mount -o prjquota /dev/$WORKER_DATA_DISK_NAME /mnt/disks/$WORKER_DATA_DISK_NAME
sudo chmod a+w /mnt/disks/$WORKER_DATA_DISK_NAME
XFS_DEVICE=$(xfs_info /mnt/disks/$WORKER_DATA_DISK_NAME | head -n 1 | awk '{{ print $1 }}' | awk  'BEGIN {{ FS = "=" }}; {{ print $2 }}')

# reconfigure docker to use local SSD
sudo service docker stop
sudo mv /var/lib/docker /mnt/disks/$WORKER_DATA_DISK_NAME/docker
sudo ln -s /mnt/disks/$WORKER_DATA_DISK_NAME/docker /var/lib/docker
sudo service docker start
<commit_msg>[batch] restore internal network (#9903)<commit_after>import os
import logging
import base64
import json

from hailtop import aiogoogle

from ..batch_configuration import PROJECT, DEFAULT_NAMESPACE
from ..worker_config import WorkerConfig
from ..log_store import LogStore

log = logging.getLogger('create_instance')

BATCH_WORKER_IMAGE = os.environ['HAIL_BATCH_WORKER_IMAGE']

log.info(f'BATCH_WORKER_IMAGE {BATCH_WORKER_IMAGE}')


async def create_instance(app, zone, machine_name, machine_type, activation_token,
                          max_idle_time_msecs, worker_local_ssd_data_disk,
                          worker_pd_ssd_data_disk_size_gb, boot_disk_size_gb):
    log_store: LogStore = app['log_store']
    compute_client: aiogoogle.ComputeClient = app['compute_client']

    if worker_local_ssd_data_disk:
        worker_data_disk = {
            'type': 'SCRATCH',
            'autoDelete': True,
            'interface': 'NVME',
            'initializeParams': {
                'diskType': f'zones/{zone}/diskTypes/local-ssd'
            }}
        worker_data_disk_name = 'nvme0n1'
    else:
        worker_data_disk = {
            'autoDelete': True,
            'initializeParams': {
                'diskType': f'projects/{PROJECT}/zones/{zone}/diskTypes/pd-ssd',
                'diskSizeGb': str(worker_pd_ssd_data_disk_size_gb)
            }
        }
        worker_data_disk_name = 'sdb'

    config = {
        'name': machine_name,
        'machineType': f'projects/{PROJECT}/zones/{zone}/machineTypes/{machine_type}',
        'labels': {
            'role': 'batch2-agent',
            'namespace': DEFAULT_NAMESPACE
        },

        'disks': [{
            'boot': True,
            'autoDelete': True,
            'initializeParams': {
                'sourceImage': f'projects/{PROJECT}/global/images/batch-worker-12',
                'diskType': f'projects/{PROJECT}/zones/{zone}/diskTypes/pd-ssd',
                'diskSizeGb': str(boot_disk_size_gb)
            }
        }, worker_data_disk],

        'networkInterfaces': [{
            'network': 'global/networks/default',
            'networkTier': 'PREMIUM',
            'accessConfigs': [{
                'type': 'ONE_TO_ONE_NAT',
                'name': 'external-nat'
            }]
        }],

        'scheduling': {
            'automaticRestart': False,
            'onHostMaintenance': "TERMINATE",
            'preemptible': True
        },

        'serviceAccounts': [{
            'email': f'<EMAIL>',
            'scopes': [
                'https://www.googleapis.com/auth/cloud-platform'
            ]
        }],

        'metadata': {
            'items': [{
                'key': 'startup-script',
                'value': '''
#!/bin/bash
set -x

curl -s -H "Metadata-Flavor: Google" "http://metadata.google.internal/computeMetadata/v1/instance/attributes/run_script"  >./run.sh

nohup /bin/bash run.sh >run.log 2>&1 &
    '''
            }, {
                'key': 'run_script',
                'value': rf'''
#!/bin/bash
set -x

# only allow udp/53 (dns) to metadata server
# -I inserts at the head of the chain, so the ACCEPT rule runs first
iptables -I DOCKER-USER -d 169.254.169.254 -j DROP
iptables -I DOCKER-USER -d 169.254.169.254 -p udp -m udp --destination-port 53 -j ACCEPT

docker network create public --opt com.docker.network.bridge.name=public
docker network create private --opt com.docker.network.bridge.name=private
# make other docker containers not route-able
iptables -I DOCKER-USER -i public -d 172.16.0.0/12 -j DROP
# not used, but ban it anyway!
iptables -I DOCKER-USER -i public -d 192.168.0.0/16 -j DROP

WORKER_DATA_DISK_NAME="{worker_data_disk_name}"

# format local SSD
sudo mkfs.xfs -m reflink=1 /dev/$WORKER_DATA_DISK_NAME
sudo mkdir -p /mnt/disks/$WORKER_DATA_DISK_NAME
sudo mount -o prjquota /dev/$WORKER_DATA_DISK_NAME /mnt/disks/$WORKER_DATA_DISK_NAME
sudo chmod a+w /mnt/disks/$WORKER_DATA_DISK_NAME
XFS_DEVICE=$(xfs_info /mnt/disks/$WORKER_DATA_DISK_NAME | head -n 1 | awk '{{ print $1 }}' | awk  'BEGIN {{ FS = "=" }}; {{ print $2 }}')

# reconfigure docker to use local SSD
sudo service docker stop
sudo mv /var/lib/docker /mnt/disks/$WORKER_DATA_DISK_NAME/docker
sudo ln -s /mnt/disks/$WORKER_DATA_DISK_NAME/docker /var/lib/docker
sudo service docker start

{'fraction_non_alphanumeric': 0.12655315232397607, 'fraction_numerical': 0.015186378278877128, 'mean_word_length': 2.383806928766057, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 11, 'https://': 2, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 9, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4424248', 'n_tokens_mistral': 3013, 'n_tokens_neox': 2833, 'n_words': 593}