Skip to content

Commit

Permalink
s
Browse files Browse the repository at this point in the history
  • Loading branch information
ayemos committed Mar 11, 2019
1 parent 4abf66e commit 91fe281
Show file tree
Hide file tree
Showing 17 changed files with 314 additions and 158 deletions.
153 changes: 29 additions & 124 deletions babysage/cli.py
Original file line number Diff line number Diff line change
@@ -1,142 +1,51 @@
import os
import boto3
import click
import pkg_resources

from jinja2 import Template

from babysage.config import get_config


@click.group()
@click.option('--debug/--no-debug', default=False)
def main(debug):
click.echo('Debug mode is %s' % ('on' if debug else 'off'))
pass


@main.command() # @cli, not @click!
@click.option('--name', nargs=1, type=str, required=True)
@click.option('--ecr-repo-name', '-r', 'ecr_repo_name', nargs=1, type=str)
@click.option('--s3-bucket-name', '-b', 's3_bucket_name', nargs=1, type=str)
@click.option('--execution-role-arn', '-e', 'execution_role_arn', nargs=1, type=str)
@click.option('--region', '-r', 'region', nargs=1, type=str)
def init(name, ecr_repo_name, s3_bucket_name, region):
click.echo(f'Creating babysage experiment {name}.')

region_name = region or boto3.session.Session().region_name
ecr = boto3.client('ecr')
ecr_repo_name = ecr_repo_name or f'babysage-{name}'

try:
resp = ecr.describe_repositories(repositoryNames=[ecr_repo_name])
ecr_repo_arn = resp['repositories'][0]['repositoryArn']
except Exception as e:
print(e.__dict__)
if e.response['Error']['Code'] == 'RepositoryNotFoundException':
click.echo(f'ECR repositroy {ecr_repo_name} doesn\'t seem to be exist')
if input('Can I create one? [Y/n]') in 'yY ':
print('create')
pass
# resp = ecr.create_repository(repositoryName=ecr_repo_name)
# ecr_repo_arn = resp['repository']['repositoryArn']
ecr_repo_arn = 'dummy'
else:
print('not create')
ecr_repo_arn = '$ECR_REPO_ARN$'
pass

s3 = boto3.client('s3')
s3_bucket_name = s3_bucket_name or f'babysage-{region_name}'

try:
s3.get_bucket_location(Bucket=s3_bucket_name)
except Exception as e:
if e.response['Error']['Code'] == 'NoSuchBucket':
# bucket doesn't exist
click.echo(f'S3 bucket {s3_bucket_name} doesn\'t seem to be exist')
if input('Can I create one? [Y/n]') in 'yY ':
print('create')
pass
# s3.create_bucket(Bucket=s3_bucket_name)
else:
print('not create')
pass

template_dir = 'templates/default'

ctx = {
'babysage':
{
'experiment_name': name,
'ecr_repo_arn': ecr_repo_arn,
's3_bucket_name': s3_bucket_name}}

for t in _expand_template(template_dir):
out_path = Template(os.path.join(t[len(template_dir) + 1:])).render(ctx)
if not os.path.isdir(os.path.dirname(out_path)):
os.makedirs(os.path.dirname(out_path))

with open(pkg_resources.resource_filename('babysage', t), 'r') as f_in:
with open(out_path, 'w') as f_out:
f_out.write(
Template(f_in.read()).render(ctx))


def _expand_template(root):
if pkg_resources.resource_isdir('babysage', root):
resources = []
for r in pkg_resources.resource_listdir('babysage', root):
resources += _expand_template(os.path.join(root, r))
return resources
else:
if pkg_resources.resource_exists('babysage', root):
return [root]
else:
raise Exception(f'Couldn\'t find the resource {root} on package `babysage`')
def init(
name,
ecr_repo_name,
s3_bucket_name,
execution_role_arn,
region):
from babysage.commands.init import init
init(
name,
ecr_repo_name,
s3_bucket_name,
execution_role_arn,
region)


@main.command()
def run_local():
click.echo('まだだよ')
@click.option('--docker-file', '-f', 'dockerfile', nargs=1, type=str, default=None)
def run_local(dockerfile):
from babysage.commands.run_local import run_local
run_local(dockerfile)


@main.command()
@click.option('--docker-file', '-f', 'dockerfile', nargs=1, type=str, default=None)
def deploy(dockerfile):
from babysage.commands.deploy import deploy
deploy(dockerfile)


@main.command()
def run_remote():
training_params = {
"AlgorithmSpecification": {
"TrainingImage": '873096238884.dkr.ecr.us-west-2.amazonaws.com/konpeki/ml-car:00300_initial_experiment',
"TrainingInputMode": "File"
},
"RoleArn": 'arn:aws:iam::873096238884:role/service-role/AmazonSageMaker-ExecutionRole-20180516T165044',
"OutputDataConfig": {
"S3OutputPath": 's3://aibs-oregon-ml-test-data/ml-car/output/00300_initial_experiment/'
},
"ResourceConfig": {
"InstanceCount": 1,
"InstanceType": "ml.p3.2xlarge",
"VolumeSizeInGB": 300
},
"TrainingJobName": job_name,
# "HyperParameters": {
# "top_k": str(top_k),
# },
"InputDataConfig": [
{
"ChannelName": "main",
"DataSource": {
"S3DataSource": {
"S3DataType": "S3Prefix",
"S3Uri": 's3://aibs-oregon-ml-test-data/ml-car/data/00300_initial_experiment/',
"S3DataDistributionType": "FullyReplicated"
}
},
"ContentType": "application/x-image",
"CompressionType": "None"
}
],
"StoppingCondition": {
"MaxRuntimeInSeconds": 360000
},
}
from babysage.commands.run_remote import run_remote
run_remote()


@main.command()
Expand All @@ -150,11 +59,7 @@ def data_down():


@main.command()
@main.argument('hash_hint')
@click.argument('hash_hint')
def logs(hash_hint):
click.echo('まだだよ')
click.echo(hash_hint)


# resource_stream(package_or_requirement, resource_name)
# resource_string(package_or_requirement, resource_name)
42 changes: 42 additions & 0 deletions babysage/commands/deploy.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,42 @@
import base64
from urllib.parse import urlparse

import boto3
import docker

from babysage.utils import find_project_root
from babysage.config import get_config


def deploy(dockerfile):
aws_account_id = boto3.client('sts').get_caller_identity()['Account']
ecr = boto3.client(service_name='ecr')
resp = ecr.get_authorization_token(
registryIds=[aws_account_id])
token = resp['authorizationData'][0]['authorizationToken']
proxy_endpoint = resp['authorizationData'][0]['proxyEndpoint']
tag = f"babysage-{get_config('experiment_name')}"
repository_tag = urlparse(proxy_endpoint).netloc + '/' + tag
user, password = base64.b64decode(token).decode().split(':')

client = docker.from_env()
client.login(
username=user,
password=password,
registry=proxy_endpoint)

project_root = find_project_root()

if project_root:
dockerfile = dockerfile or 'Dockerfile'
client.images.build(
path=project_root,
quiet=False,
rm=True,
tag=repository_tag,
dockerfile=dockerfile)
for l in client.images.push(repository_tag, stream=True, decode=True):
print(l)
else:
raise Exception('''
Current directry is not a babysage project directry (could not find `babysage.yml` file''')
90 changes: 90 additions & 0 deletions babysage/commands/init.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,90 @@
import os
import pkg_resources

import boto3
from jinja2 import Template


def init(name, ecr_repo_name, s3_bucket_name, execution_role_arn, region):
aws_account_id = boto3.client('sts').get_caller_identity()['Account']

print(f'Creating babysage experiment {name}.')

session = boto3.session.Session()
region_name = region or session.region_name
ecr = session.client('ecr')
ecr_repo_name = ecr_repo_name or f'babysage-{name}'
ecr_repo_arn = f'{aws_account_id}.dkr.ecr.{region_name}.amazonaws.com/{ecr_repo_name}'

try:
resp = ecr.describe_repositories(repositoryNames=[ecr_repo_name])
ecr_repo_arn = resp['repositories'][0]['repositoryArn']
except Exception as e:
print(e.__dict__)
if e.response['Error']['Code'] == 'RepositoryNotFoundException':
print(f'ECR repositroy {ecr_repo_arn} doesn\'t seem to be exist')
if input('Can I create one? [y/N]') in 'yY':
resp = ecr.create_repository(repositoryName=ecr_repo_name)

s3 = session.client('s3')
s3_bucket_name = s3_bucket_name or f'babysage-{aws_account_id}-{region_name}'

try:
s3.get_bucket_location(Bucket=s3_bucket_name)
except Exception as e:
if e.response['Error']['Code'] == 'NoSuchBucket':
# bucket doesn't exist
print(f'S3 bucket {s3_bucket_name} doesn\'t seem to be exist')
if input('Can I create one? [y/N]') in 'yY':
s3.create_bucket(
Bucket=s3_bucket_name,
CreateBucketConfiguration={'LocationConstraint': region_name})

'''
iam = session.client('iam')
try:
iam.get_bucket_location(Bucket=s3_bucket_name)
except Exception as e:
if e.response['Error']['Code'] == 'NoSuchBucket':
# bucket doesn't exist
click.echo(f'S3 bucket {s3_bucket_name} doesn\'t seem to be exist')
if input('Can I create one? [Y/n]') in 'yY ':
s3.create_bucket(
Bucket=s3_bucket_name,
CreateBucketConfiguration={'LocationConstraint': region_name})
'''

template_dir = 'templates/default'

ctx = {
'babysage':
{
'experiment_name': name,
'ecr_repo_name': ecr_repo_name,
's3_bucket_name': s3_bucket_name}}

for t in _expand_template(template_dir):
out_path = Template(os.path.join(t[len(template_dir) + 1:])).render(ctx)
if not os.path.isdir(os.path.dirname(out_path)):
os.makedirs(os.path.dirname(out_path))

with open(pkg_resources.resource_filename('babysage', t), 'r') as f_in:
with open(out_path, 'w') as f_out:
f_out.write(
Template(f_in.read()).render(ctx))
if os.path.basename(out_path) == 'train':
import stat
os.chmod(out_path, stat.S_IEXEC)


def _expand_template(root):
if pkg_resources.resource_isdir('babysage', root):
resources = []
for r in pkg_resources.resource_listdir('babysage', root):
resources += _expand_template(os.path.join(root, r))
return resources
else:
if pkg_resources.resource_exists('babysage', root):
return [root]
else:
raise Exception(f'Couldn\'t find the resource {root} on package `babysage`')
33 changes: 33 additions & 0 deletions babysage/commands/run_local.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,33 @@
import os

import docker

from babysage.utils import find_project_root, get_input_dir, get_output_dir
from babysage.config import get_config


def run_local(dockerfile):
client = docker.from_env()

project_root = find_project_root()

if project_root:
dockerfile = dockerfile or 'Dockerfile'
tag = f"babysage-{get_config('experiment_name')}",
client.images.build(
path=project_root,
quiet=False,
rm=True,
tag=tag,
dockerfile=dockerfile)
client.containers.run(
tag,
'train',
volumes={
os.path.join(project_root, 'input'): {
'bind': get_input_dir(), 'mode': 'rw'},
os.path.join(project_root, 'output'): {
'bind': get_output_dir(), 'mode': 'rw'}})
else:
raise Exception('''
Current directry is not a babysage project directry (could not find `babysage.yml` file''')
Loading

0 comments on commit 91fe281

Please sign in to comment.