entry_point.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp entry_point"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "\n",
    "import os\n",
    "import argparse\n",
    "import time\n",
    "import importlib\n",
    "import inspect\n",
    "import logging\n",
    "import re\n",
    "from gcp_runner.ai_platform_constants import DistributionStrategyType\n",
    "\n",
    "def get_distribution_strategy_instance(distribution_strategy_type):\n",
    "    import tensorflow as tf\n",
    "\n",
    "    logging.info('initializing distribution strategy: %s', distribution_strategy_type)\n",
    "    if distribution_strategy_type == DistributionStrategyType.TPU_STRATEGY:\n",
    "        tpu = None\n",
    "        try:\n",
    "            logging.info('resolving to TPU cluster')\n",
    "            tpu = tf.distribute.cluster_resolver.TPUClusterResolver()\n",
    "            logging.info('connecting to TPU cluster')\n",
    "            tf.config.experimental_connect_to_cluster(tpu)\n",
    "        except ValueError as e:\n",
    "            logging.info('error connecting to TPU cluster: %s', e)\n",
    "            return None\n",
    "        logging.info('initializing TPU system')\n",
    "        tf.tpu.experimental.initialize_tpu_system(tpu)\n",
    "        distribution_strategy = tf.distribute.experimental.TPUStrategy(tpu)\n",
    "        logging.info(\n",
    "            'training using TPUStrategy, tpu.cluster_spec: %s', tpu.cluster_spec())\n",
    "        return distribution_strategy\n",
    "    elif distribution_strategy_type == DistributionStrategyType.ONE_DEVICE_STRATEGY:\n",
    "        return tf.distribute.OneDeviceStrategy(device=\"/cpu:0\")\n",
    "    elif distribution_strategy_type:\n",
    "        return eval(distribution_strategy_type.value)()\n",
    "    \n",
    "def parse_unknown_args(unknown_args):\n",
    "    kwargs = {}\n",
    "    for arg in unknown_args:\n",
    "        parse_result = re.search('--(.+)=(.+)', arg, re.IGNORECASE)\n",
    "        if parse_result is not None and parse_result.groups() is not None and len(parse_result.groups()) == 2:\n",
    "            kwargs[parse_result.group(1).replace('-', '_')] = parse_result.group(2)\n",
    "        else:\n",
    "            print(\"can't parse argument: %s\" % arg)\n",
    "    return kwargs\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "get_distribution_strategy_instance(None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "\n",
    "import os\n",
    "import json\n",
    "import socket\n",
    "\n",
    "def setup_keras_tuner_config():\n",
    "    if 'TF_CONFIG' in os.environ:\n",
    "        try:\n",
    "            tf_config = json.loads(os.environ['TF_CONFIG'])\n",
    "            cluster = tf_config['cluster']\n",
    "            task = tf_config['task']\n",
    "            chief_addr = cluster['chief'][0].split(':')\n",
    "            chief_ip = socket.gethostbyname(chief_addr[0])\n",
    "            chief_port = chief_addr[1]\n",
    "            os.environ['KERASTUNER_ORACLE_IP'] = chief_ip\n",
    "            os.environ['KERASTUNER_ORACLE_PORT'] = chief_port\n",
    "            if task['type'] == 'chief':\n",
    "                os.environ['KERASTUNER_TUNER_ID'] = 'chief'\n",
    "            else:\n",
    "                os.environ['KERASTUNER_TUNER_ID'] = 'tuner{}'.format(task['index'])\n",
    "\n",
    "            print('set following environment arguments:')\n",
    "            print('KERASTUNER_ORACLE_IP: %s' % os.environ['KERASTUNER_ORACLE_IP'])\n",
    "            print('KERASTUNER_ORACLE_PORT: %s' % os.environ['KERASTUNER_ORACLE_PORT'])\n",
    "            print('KERASTUNER_TUNER_ID: %s' % os.environ['KERASTUNER_TUNER_ID'])\n",
    "        except Exception as ex:\n",
    "            print('Error setting up keras tuner config: %s' % str(ex))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "import sys\n",
    "\n",
    "def main():\n",
    "    print('in gcp_runner entry point:')\n",
    "    print(sys.argv)\n",
    "    args_parser = argparse.ArgumentParser()\n",
    "    args_parser.add_argument(\n",
    "        '--module-name',\n",
    "        help='module name of an app to run',\n",
    "        required=True)\n",
    "    args_parser.add_argument(\n",
    "        '--function-name',\n",
    "        help='function name to run',\n",
    "        required=True)\n",
    "    args_parser.add_argument(\n",
    "        '--distribution-strategy-type',\n",
    "        help='distribution strategy',\n",
    "        choices=list([e.value for e in DistributionStrategyType]))\n",
    "    args_parser.add_argument(\n",
    "        '--use-distribution-strategy-scope',\n",
    "        action='store_true',\n",
    "        help='whether to run training in a distribution strategy scope',\n",
    "        default=False)\n",
    "\n",
    "    setup_keras_tuner_config()    \n",
    "    args, unknown_args = args_parser.parse_known_args()\n",
    "    print('running entrypoint function: %s.%s' % (args.module_name, args.function_name))\n",
    "    module = importlib.import_module(args.module_name)\n",
    "    func = getattr(module, args.function_name)\n",
    "    if unknown_args is None or len(unknown_args) == 0:\n",
    "        return func()\n",
    "    \n",
    "    print('additional args: %s' % str(unknown_args))\n",
    "    args_spec = inspect.getargspec(func)\n",
    "    print(args_spec)\n",
    "    if args_spec is None or args_spec.keywords is None or len(args_spec.keywords) == 0:\n",
    "        print('provided function does not take any arguments, running as is')\n",
    "        return func()\n",
    "        \n",
    "    distribution_strategy = None\n",
    "    kwargs = parse_unknown_args(unknown_args)\n",
    "    if (args.distribution_strategy_type is not None):\n",
    "        distribution_strategy_type = DistributionStrategyType(args.distribution_strategy_type)\n",
    "        distribution_strategy = get_distribution_strategy_instance(distribution_strategy_type)\n",
    "        kwargs['distribution_strategy_type'] = distribution_strategy_type\n",
    "        \n",
    "    if distribution_strategy is not None:\n",
    "        kwargs['distribution_strategy'] = distribution_strategy\n",
    "        if args.use_distribution_strategy_scope:\n",
    "            print('running code in %s scope' % args.distribution_strategy_type)\n",
    "            with distribution_strategy.scope():\n",
    "                return func(**kwargs)\n",
    "    \n",
    "    func(**kwargs)\n",
    "    \n",
    "if __name__ == '__main__':\n",
    "    main()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "\n",
    "from gcp_runner import core\n",
    "core.export_and_reload_all(silent=True, ignore_errors=False)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.5 64-bit",
   "language": "python",
   "name": "python38564bit13e786327e864b6b8a0b25f674ce5abb"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}