<commit_before># Copyright 2016 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.

"""Methods for implementing the `datalab create` command."""

import os
import subprocess
import tempfile

import connect
import utils


description = ("""`{0} {1}` creates a new Datalab instances running in a Google
Compute Engine VM.

This command also creates the 'datalab-network' network if necessary.

By default, the command creates a persistent connection to the newly
created instance. You can disable that behavior by passing in the
'--no-connect' flag.""")


_DATALAB_NETWORK = 'datalab-network'
_DATALAB_NETWORK_DESCRIPTION = 'Network for Google Cloud Datalab instances'

_DATALAB_FIREWALL_RULE = 'datalab-network-allow-ssh'
_DATALAB_FIREWALL_RULE_DESCRIPTION = 'Allow SSH access to Datalab instances'

_DATALAB_DEFAULT_DISK_SIZE_GB = 200
_DATALAB_DISK_DESCRIPTION = (
    'Persistent disk for a Google Cloud Datalab instance')

_DATALAB_NOTEBOOKS_REPOSITORY = 'datalab-notebooks'

_DATALAB_STARTUP_SCRIPT = """#!/bin/bash

PERSISTENT_DISK_DEV="/dev/disk/by-id/google-datalab-pd"
MOUNT_DIR="/mnt/disks/datalab-pd"
MOUNT_CMD="mount -o discard,defaults ${{PERSISTENT_DISK_DEV}} ${{MOUNT_DIR}}"

clone_repo() {{
  echo "Creating the datalab directory"
  mkdir -p ${{MOUNT_DIR}}/content/datalab
  echo "Cloning the repo {0}"
  docker run --rm -v "${{MOUNT_DIR}}/content:/content" \
    --entrypoint "/bin/bash" {0} \
    gcloud source repos clone {1} /content/datalab/notebooks
}}

repo_is_populated() {{
  cd ${{MOUNT_DIR}}/content/datalab/notebooks
  git show-ref --quiet
}}

populate_repo() {{
  echo "Populating datalab-notebooks repo"
  docker run --rm -v "${{MOUNT_DIR}}/content:/content" \
    --workdir=/content/datalab/notebooks \
    --entrypoint "/bin/bash" {0} -c "\
        echo '.ipynb_checkpoints' >> .gitignore; \
        echo '*.pyc' >> .gitignore; \
        echo '# Project Notebooks' >> README.md; \
        git add .gitignore README.md; \
        git -c user.email=nobody -c user.name=Datalab \
          commit --message='Set up initial notebook repo.'; \
        git push origin master; \
    "
}}

format_disk() {{
  echo "Formatting the persistent disk"
  mkfs.ext4 -F \
    -E lazy_itable_init=0,lazy_journal_init=0,discard \
    ${{PERSISTENT_DISK_DEV}}
  ${{MOUNT_CMD}}
  clone_repo
  if ! repo_is_populated; then
    populate_repo
  fi
}}

mount_and_prepare_disk() {{
  echo "Trying to mount the persistent disk"
  mkdir -p "${{MOUNT_DIR}}"
  ${{MOUNT_CMD}} || format_disk
  chmod a+w "${{MOUNT_DIR}}"
  mkdir -p "${{MOUNT_DIR}}/content"

  old_dir="${{MOUNT_DIR}}/datalab"
  new_dir="${{MOUNT_DIR}}/content/datalab"
  if [ -d "${{old_dir}}" ] && [ ! -d "${{new_dir}}" ]; then
    echo "Moving ${{old_dir}} to ${{new_dir}}"
    mv "${{old_dir}}" "${{new_dir}}"
  else
    echo "Creating ${{new_dir}}"
    mkdir -p "${{new_dir}}"
  fi
}}

configure_swap() {{
  mem_total_line=`cat /proc/meminfo | grep MemTotal`
  mem_total_value=`echo "${{mem_total_line}}" | cut -d ':' -f 2`
  memory_kb=`echo "${{mem_total_value}}" | cut -d 'k' -f 1 | tr -d '[:space:]'`
  swapfile="${{MOUNT_DIR}}/swapfile"

  # Create the swapfile if it is either missing or not big enough
  current_size="0"
  if [ -e "${{swapfile}}" ]; then
    current_size=`ls -s ${{swapfile}} | cut -d ' ' -f 1`
  fi
  if [ "${{memory_kb}}" -gt "${{current_size}}" ]; then
    echo "Creating a ${{memory_kb}} kilobyte swapfile at ${{swapfile}}"
    dd if=/dev/zero of="${{swapfile}}" bs=1024 count="${{memory_kb}}"
  fi
  chmod 0600 "${{swapfile}}"
  mkswap "${{swapfile}}"

  # Enable swap
  sysctl vm.disk_based_swap=1
  swapon "${{swapfile}}"
}}

cleanup_tmp() {{
  tmpdir="${{MOUNT_DIR}}/tmp"
  rm -rf "${{tmpdir}}"
  mkdir -p "${{tmpdir}}"
}}

mount_and_prepare_disk
configure_swap
cleanup_tmp

journalctl -u google-startup-scripts --no-pager > /var/log/startupscript.log
"""

_DATALAB_CONTAINER_MANIFEST_URL = (
    'http://metadata.google.internal/' +
    'computeMetadata/v1/instance/attributes/google-container-manifest')

_DATALAB_CLOUD_CONFIG = """
#cloud-config

runcmd:
- ['curl', '-X', 'GET', '-H', 'Metadata-Flavor: Google','{0}',
   '-o', '/tmp/podspec.yaml']
- ['kubelet', '--pod-manifest-path', '/tmp/podspec.yaml']
""".format(_DATALAB_CONTAINER_MANIFEST_URL)

_DATALAB_CONTAINER_SPEC = """
apiVersion: v1
kind: Pod
metadata:
  name: 'datalab-server'
spec:
  containers:
    - name: datalab<commit_msg>Two fixes for the race condition in setting up the /tmp dir: (#1296)<commit_after># Copyright 2016 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.

"""Methods for implementing the `datalab create` command."""

import os
import subprocess
import tempfile

import connect
import utils


description = ("""`{0} {1}` creates a new Datalab instances running in a Google
Compute Engine VM.

This command also creates the 'datalab-network' network if necessary.

By default, the command creates a persistent connection to the newly
created instance. You can disable that behavior by passing in the
'--no-connect' flag.""")


_DATALAB_NETWORK = 'datalab-network'
_DATALAB_NETWORK_DESCRIPTION = 'Network for Google Cloud Datalab instances'

_DATALAB_FIREWALL_RULE = 'datalab-network-allow-ssh'
_DATALAB_FIREWALL_RULE_DESCRIPTION = 'Allow SSH access to Datalab instances'

_DATALAB_DEFAULT_DISK_SIZE_GB = 200
_DATALAB_DISK_DESCRIPTION = (
    'Persistent disk for a Google Cloud Datalab instance')

_DATALAB_NOTEBOOKS_REPOSITORY = 'datalab-notebooks'

_DATALAB_STARTUP_SCRIPT = """#!/bin/bash

PERSISTENT_DISK_DEV="/dev/disk/by-id/google-datalab-pd"
MOUNT_DIR="/mnt/disks/datalab-pd"
MOUNT_CMD="mount -o discard,defaults ${{PERSISTENT_DISK_DEV}} ${{MOUNT_DIR}}"

clone_repo() {{
  echo "Creating the datalab directory"
  mkdir -p ${{MOUNT_DIR}}/content/datalab
  echo "Cloning the repo {0}"
  docker run --rm -v "${{MOUNT_DIR}}/content:/content" \
    --entrypoint "/bin/bash" {0} \
    gcloud source repos clone {1} /content/datalab/notebooks
}}

repo_is_populated() {{
  cd ${{MOUNT_DIR}}/content/datalab/notebooks
  git show-ref --quiet
}}

populate_repo() {{
  echo "Populating datalab-notebooks repo"
  docker run --rm -v "${{MOUNT_DIR}}/content:/content" \
    --workdir=/content/datalab/notebooks \
    --entrypoint "/bin/bash" {0} -c "\
        echo '.ipynb_checkpoints' >> .gitignore; \
        echo '*.pyc' >> .gitignore; \
        echo '# Project Notebooks' >> README.md; \
        git add .gitignore README.md; \
        git -c user.email=nobody -c user.name=Datalab \
          commit --message='Set up initial notebook repo.'; \
        git push origin master; \
    "
}}

format_disk() {{
  echo "Formatting the persistent disk"
  mkfs.ext4 -F \
    -E lazy_itable_init=0,lazy_journal_init=0,discard \
    ${{PERSISTENT_DISK_DEV}}
  ${{MOUNT_CMD}}
  clone_repo
  if ! repo_is_populated; then
    populate_repo
  fi
}}

mount_and_prepare_disk() {{
  echo "Trying to mount the persistent disk"
  mkdir -p "${{MOUNT_DIR}}"
  ${{MOUNT_CMD}} || format_disk
  chmod a+w "${{MOUNT_DIR}}"
  mkdir -p "${{MOUNT_DIR}}/content"

  old_dir="${{MOUNT_DIR}}/datalab"
  new_dir="${{MOUNT_DIR}}/content/datalab"
  if [ -d "${{old_dir}}" ] && [ ! -d "${{new_dir}}" ]; then
    echo "Moving ${{old_dir}} to ${{new_dir}}"
    mv "${{old_dir}}" "${{new_dir}}"
  else
    echo "Creating ${{new_dir}}"
    mkdir -p "${{new_dir}}"
  fi
}}

configure_swap() {{
  mem_total_line=`cat /proc/meminfo | grep MemTotal`
  mem_total_value=`echo "${{mem_total_line}}" | cut -d ':' -f 2`
  memory_kb=`echo "${{mem_total_value}}" | cut -d 'k' -f 1 | tr -d '[:space:]'`
  swapfile="${{MOUNT_DIR}}/swapfile"

  # Create the swapfile if it is either missing or not big enough
  current_size="0"
  if [ -e "${{swapfile}}" ]; then
    current_size=`ls -s ${{swapfile}} | cut -d ' ' -f 1`
  fi
  if [ "${{memory_kb}}" -gt "${{current_size}}" ]; then
    echo "Creating a ${{memory_kb}} kilobyte swapfile at ${{swapfile}}"
    dd if=/dev/zero of="${{swapfile}}" bs=1024 count="${{memory_kb}}"
  fi
  chmod 0600 "${{swapfile}}"
  mkswap "${{swapfile}}"

  # Enable swap
  sysctl vm.disk_based_swap=1
  swapon "${{swapfile}}"
}}

cleanup_tmp() {{
  tmpdir="${{MOUNT_DIR}}/tmp"

  # First, make sure the temporary directory exists.
  mkdir -p "${{tmpdir}}"

  # Remove all files from it.
  #
  # We do not remove the directory itself, as that could lead to a broken
  # volume mount if the Docker container has already started).
  #
  # We also do not just use `rm -rf ${{tmpdir}}/*`, as that would leave
  # behind any hidden files.
  find "${{tmpdir}}/" -mindepth 1 -delete
}}

mount_and_prepare_disk
configure_swap
cleanup_tmp

journalctl -u google-startup-scripts --no-pager > /var/log/startupscript.log
"""

_DATALAB_CONTAINER_MANIFEST_URL = (
    'http://metadata.google.internal/' +
    'computeMetadata/v1/instance/attributes/google-container-manifest')

_DATALAB_CLOUD_CONFIG = """
#cloud-config

runcmd:
- ['while', '[', '!', '-e', '/mnt/disks/datalab-pd/tmp', ']', ';',
   'do', 'sleep', '1', ';', 'done']
- ['curl', '-X', 'GET', '-H', 'Metadata-Flavor: Google','{0}',
   '-o', '/tmp/podspec.yaml']
- ['kubelet', '--pod-manifest-path', '/tmp/podspec.yaml']
""".format(_DATALAB_CONTAINER_MANIFEST_URL)

_DATALAB_CONTAINER_SPEC = """
apiVersion: v1
kind: Pod
metadata:
  name: 'datalab-server'
spec:
  containers:
    - name: datalab
{'fraction_non_alphanumeric': 0.1549159745026077, 'fraction_numerical': 0.007533320455862469, 'mean_word_length': 3.985556090515166, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 17, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, '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': '1979232', 'n_tokens_mistral': 3825, 'n_tokens_neox': 3661, 'n_words': 978}