# Path: studio/git_util.py
# def get_git_info(path='.', abort_dirty=True):
# def is_git(path='.'):
# def is_clean(path='.'):
# def get_repo_url(path='.', remove_user=True):
# def get_branch(path='.'):
# def get_commit(path='.'):
# def get_my_repo_url():
# def get_my_branch():
# def get_my_checkout_target():
# 
# Path: studio/util/logs.py
# DEBUG = logging.DEBUG
# INFO = logging.INFO
# ERROR = logging.ERROR
# def get_logger(name):
# def debug(line):
# def error(line):
# def info(line):
# 
# Path: studio/util/gpu_util.py
# def memstr2int(string):
#     conversion_factors = [
#         ('Mb', 2**20), ('MiB', 2**20), ('m', 2**20), ('mb', 2**20),
#         ('Gb', 2**30), ('GiB', 2**30), ('g', 2**30), ('gb', 2**30),
#         ('kb', 2**10), ('k', 2**10)
#     ]
# 
#     for key, factor in conversion_factors:
#         if string.endswith(key):
#             return int(float(string.replace(key, '')) * factor)
# 
#     return int(string)
# 
# Path: studio/cloud_worker_util.py
# def insert_user_startup_script(user_startup_script, startup_script_str,
#                                logger):
#     if user_startup_script is None:
#         return startup_script_str
# 
#     try:
#         with open(os.path.abspath(os.path.expanduser(
#                 user_startup_script))) as f:
#             user_startup_script_lines = f.read().splitlines()
#     except BaseException:
#         if user_startup_script is not None:
#             logger.warn("User startup script (%s) cannot be loaded" %
#                         user_startup_script)
#         return startup_script_str
# 
#     startup_script_lines = startup_script_str.splitlines()
#     new_startup_script_lines = []
#     whitespace = " " * INDENT
#     for line in startup_script_lines:
# 
#         if line.startswith("studio remote worker") or \
#                 line.startswith("studio-remote-worker"):
#             curr_working_dir = "curr_working_dir_%s" % rand_string(32)
#             func_name = "user_script_%s" % rand_string(32)
# 
#             new_startup_script_lines.append("%s=$(pwd)\n" % curr_working_dir)
#             new_startup_script_lines.append("cd ~\n")
#             new_startup_script_lines.append("%s()(\n" % func_name)
#             for user_line in user_startup_script_lines:
#                 if user_line.startswith("#!"):
#                     continue
#                 new_startup_script_lines.append("%s%s\n" %
#                                                 (whitespace, user_line))
# 
#             new_startup_script_lines.append("%scd $%s\n" %
#                                             (whitespace, curr_working_dir))
#             new_startup_script_lines.append("%s%s\n" %
#                                             (whitespace, line))
#             new_startup_script_lines.append(")\n")
#             new_startup_script_lines.append("%s\n" % func_name)
#         else:
#             new_startup_script_lines.append("%s\n" % line)
# 
#     new_startup_script = "".join(new_startup_script_lines)
#     logger.info('Inserting the following user startup script'
#                 ' into the default startup script:')
#     logger.info("\n".join(user_startup_script_lines))
# 
#     # with open("/home/jason/Desktop/script.sh", 'wb') as f:
#     #     f.write(new_startup_script)
#     # sys.exit()
# 
#     return new_startup_script


import os
import time
import googleapiclient.discovery
import uuid
import math
import json
from studio import git_util
from studio.util import logs
from studio.util.gpu_util import memstr2int
from .cloud_worker_util import insert_user_startup_script
#!/usr/bin/env python




class GCloudWorkerManager(object):
    def __init__(self, zone='us-east1-c',
                 auth_cookie=None, verbose=10, branch=None,
                 user_startup_script=None):
        assert 'GOOGLE_APPLICATION_CREDENTIALS' in os.environ.keys()
        with open(os.environ['GOOGLE_APPLICATION_CREDENTIALS'], 'r') as f:
            credentials_dict = json.loads(f.read())

        self.compute = googleapiclient.discovery.build('compute', 'v1')

        self.startup_script_file = os.path.join(
            os.path.dirname(__file__),
            'scripts/gcloud_worker_startup.sh')

        self.install_studio_script = os.path.join(
            os.path.dirname(__file__),
            'scripts/install_studio.sh')

        self.zone = zone
        self.projectid = credentials_dict['project_id']
        self.logger = logs.get_logger("GCloudWorkerManager")
        self.logger.setLevel(verbose)
        self.auth_cookie = auth_cookie
        self.user_startup_script = user_startup_script
        self.repo_url = git_util.get_my_repo_url()

# Path: studio/git_util.py
# def get_git_info(path='.', abort_dirty=True):
# def is_git(path='.'):
# def is_clean(path='.'):
# def get_repo_url(path='.', remove_user=True):
# def get_branch(path='.'):
# def get_commit(path='.'):
# def get_my_repo_url():
# def get_my_branch():
# def get_my_checkout_target():
# 
# Path: studio/util/logs.py
# DEBUG = logging.DEBUG
# INFO = logging.INFO
# ERROR = logging.ERROR
# def get_logger(name):
# def debug(line):
# def error(line):
# def info(line):
# 
# Path: studio/util/gpu_util.py
# def memstr2int(string):
#     conversion_factors = [
#         ('Mb', 2**20), ('MiB', 2**20), ('m', 2**20), ('mb', 2**20),
#         ('Gb', 2**30), ('GiB', 2**30), ('g', 2**30), ('gb', 2**30),
#         ('kb', 2**10), ('k', 2**10)
#     ]
# 
#     for key, factor in conversion_factors:
#         if string.endswith(key):
#             return int(float(string.replace(key, '')) * factor)
# 
#     return int(string)
# 
# Path: studio/cloud_worker_util.py
# def insert_user_startup_script(user_startup_script, startup_script_str,
#                                logger):
#     if user_startup_script is None:
#         return startup_script_str
# 
#     try:
#         with open(os.path.abspath(os.path.expanduser(
#                 user_startup_script))) as f:
#             user_startup_script_lines = f.read().splitlines()
#     except BaseException:
#         if user_startup_script is not None:
#             logger.warn("User startup script (%s) cannot be loaded" %
#                         user_startup_script)
#         return startup_script_str
# 
#     startup_script_lines = startup_script_str.splitlines()
#     new_startup_script_lines = []
#     whitespace = " " * INDENT
#     for line in startup_script_lines:
# 
#         if line.startswith("studio remote worker") or \
#                 line.startswith("studio-remote-worker"):
#             curr_working_dir = "curr_working_dir_%s" % rand_string(32)
#             func_name = "user_script_%s" % rand_string(32)
# 
#             new_startup_script_lines.append("%s=$(pwd)\n" % curr_working_dir)
#             new_startup_script_lines.append("cd ~\n")
#             new_startup_script_lines.append("%s()(\n" % func_name)
#             for user_line in user_startup_script_lines:
#                 if user_line.startswith("#!"):
#                     continue
#                 new_startup_script_lines.append("%s%s\n" %
#                                                 (whitespace, user_line))
# 
#             new_startup_script_lines.append("%scd $%s\n" %
#                                             (whitespace, curr_working_dir))
#             new_startup_script_lines.append("%s%s\n" %
#                                             (whitespace, line))
#             new_startup_script_lines.append(")\n")
#             new_startup_script_lines.append("%s\n" % func_name)
#         else:
#             new_startup_script_lines.append("%s\n" % line)
# 
#     new_startup_script = "".join(new_startup_script_lines)
#     logger.info('Inserting the following user startup script'
#                 ' into the default startup script:')
#     logger.info("\n".join(user_startup_script_lines))
# 
#     # with open("/home/jason/Desktop/script.sh", 'wb') as f:
#     #     f.write(new_startup_script)
#     # sys.exit()
# 
#     return new_startup_script


# Path: studio/gcloud_worker.py
import os
import time
import googleapiclient.discovery
import uuid
import math
import json
from studio import git_util
from studio.util import logs
from studio.util.gpu_util import memstr2int
from .cloud_worker_util import insert_user_startup_script

#!/usr/bin/env python




class GCloudWorkerManager(object):
    def __init__(self, zone='us-east1-c',
                 auth_cookie=None, verbose=10, branch=None,
                 user_startup_script=None):
        assert 'GOOGLE_APPLICATION_CREDENTIALS' in os.environ.keys()
        with open(os.environ['GOOGLE_APPLICATION_CREDENTIALS'], 'r') as f:
            credentials_dict = json.loads(f.read())

        self.compute = googleapiclient.discovery.build('compute', 'v1')

        self.startup_script_file = os.path.join(
            os.path.dirname(__file__),
            'scripts/gcloud_worker_startup.sh')

        self.install_studio_script = os.path.join(
            os.path.dirname(__file__),
            'scripts/install_studio.sh')

        self.zone = zone
        self.projectid = credentials_dict['project_id']
        self.logger = logs.get_logger("GCloudWorkerManager")
        self.logger.setLevel(verbose)
        self.auth_cookie = auth_cookie
        self.user_startup_script = user_startup_script
        self.repo_url = git_util.get_my_repo_url()

        self.branch = branch if branch else git_util.get_my_checkout_target()