# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python [conda env:root] *
#     language: python
#     name: conda-root-py
# ---

# # Distributed training with Vertex Reduction server

# +
import os
import pprint
import sys

from google.cloud import aiplatform
# -

# ## Set environment constants

# +
PROJECT_ID = 'jk-mlops-dev'
REGION = 'us-west1'
STAGING_BUCKET = 'gs://jk-vertex-staging'

BASE_IMAGE = 'gcr.io/deeplearning-platform-release/tf2-gpu.2-5'
TRAIN_IMAGE = f'gcr.io/{PROJECT_ID}/tf_nlp_toolkit'

MNLI_TRAIN_SPLIT = 'gs://jk-vertex-demos/datasets/MNLI/mnli_train.tf_record'
MNLI_VALID_SPLIT = 'gs://jk-vertex-demos/datasets/MNLI/mnli_valid.tf_record'

BERT_HUB_URL = 'https://tfhub.dev/tensorflow/bert_en_uncased_L-24_H-1024_A-16/4'
# -

# ## Train

# ### Initialize Vertex AI SDK
#

aiplatform.init(
    project=PROJECT_ID,
    location=REGION,
    staging_bucket=STAGING_BUCKET
)

# ### Create a training script

# +
# %%writefile trainer/test.py

import os

import tensorflow as tf

strategy = tf.distribute.MultiWorkerMirroredStrategy()

if strategy.cluster_resolver:    
        task_type, task_id = (strategy.cluster_resolver.task_type,
                              strategy.cluster_resolver.task_id)
else:
        task_type, task_id =(None, None)
        
print('*** task_type: {}, task_id: {}'.format(task_type, task_id))

print(os.environ['AIP_MODEL_DIR'], os.environ['AIP_CHECKPOINT_DIR'], os.environ['AIP_TENSORBOARD_LOG_DIR'])


# -

print(os.environ['PATH'])

# ### Create a training container

# +
dockerfile = f'''
FROM {BASE_IMAGE}

RUN pip install pip install tf-models-official tensorflow-text 

WORKDIR /

# Copies the trainer code to the docker image.
COPY trainer /trainer

# Sets up the entry point to invoke the trainer.
ENTRYPOINT ["python", "-m", "trainer.train"]
'''

with open('Dockerfile', 'w') as f:
    f.write(dockerfile)
# -

# ! docker build -t {TRAIN_IMAGE} .

# ### Test the container locally

# +
replica_count = 1
machine_type = 'n1-standard-8'
accelerator_count = 2
accelerator_type = 'NVIDIA_TESLA_T4'

output_dir = 'gs://jk-vertex-demos/rs-demo'
model_dir = f'{output_dir}/model'
tfhub_cache_dir = f'{output_dir}/tfhub-cache'

strategy = 'mirrored'



params_override = [
    'task.train_data.input_path=' + MNLI_TRAIN_SPLIT,
    'task.validation_data.input_path=' + MNLI_VALID_SPLIT,
    'task.hub_module_url=' + BERT_HUB_URL,
    'runtime.num_gpus=' + str(accelerator_count),
    'runtime.distribution_strategy=' + strategy,
]

params = ','.join(params_override)


# -

# !docker run -it --rm --gpus all gcr.io/jk-mlops-dev/tf_nlp_toolkit \
# --experiment=bert/sentence_prediction \
# --mode=train_and_eval \
# --model_dir=gs://jk-vertex-staging/test \
# --config_file=trainer/glue_mnli_matched.yaml \
# --params_override={params}

# ### Push the container

# ! docker push {TRAIN_IMAGE}

def prepare_worker_pool_specs(
    image_uri,
    args,
    replica_count=1,
    machine_type='n1-standard-4',
    accelerator_count=0,
    accelerator_type='ACCELERATOR_TYPE_UNSPECIFIED'):

    if accelerator_count > 0:
        machine_spec = {
            'machine_type': machine_type,
            'accelerator_type': accelerator_type,
            'accelerator_count': accelerator_count,
        }
    else:
        machine_spec = {
            'machine_type': machine_type
        }
    
    container_spec = {
        'image_uri': image_uri,
        'args': args
    }
    
    chief_spec = {
        'replica_count': 1,
        'machine_spec': machine_spec,
        'container_spec': container_spec
    }

    worker_pool_specs = [chief_spec]
    if replica_count > 1:
        workers_spec = {
            'replica_count': replica_count - 1,
            'machine_spec': machine_spec,
            'container_spec': container_spec
        }
        worker_pool_specs.append(workers_spec)
    
    return worker_pool_specs


# +
replica_count = 2
machine_type = 'n1-standard-8'
accelerator_count = 1
accelerator_type = 'NVIDIA_TESLA_V100'
global_batch_size = 32

output_dir = 'gs://jk-vertex-demos/rs-demo'
model_dir = f'{output_dir}/model'
tfhub_cache_dir = f'{output_dir}/tfhub-cache'

#strategy = 'mirrored'
strategy = 'multi_worker_mirrored'

summary_interval: 1000
    
# Training data size 392,702 examples, 3 epochs.
train_steps: 36813
validation_interval: 6135

params_override = [
    'task.train_data.input_path=' + MNLI_TRAIN_SPLIT,
    'task.validation_data.input_path=' + MNLI_VALID_SPLIT,
    'task.train_data.global_batch_size=' + str(global_batch_size),
    'task.validation_data.global_batch_size=' + str(global_batch_size),
    'task.hub_module_url=' + BERT_HUB_URL,
    'runtime.num_gpus=' + str(accelerator_count),
    'runtime.distribution_strategy=' + strategy,
]


image_uri = TRAIN_IMAGE
args = [
    '--experiment=bert/sentence_prediction',
    '--mode=train_and_eval',
    '--model_dir=' + model_dir,
    '--config_file=trainer/glue_mnli_matched.yaml',
    '--tfhub_cache_dir=' + tfhub_cache_dir,
    '--params_override=' + ','.join(params_override),
]

worker_pool_specs = prepare_worker_pool_specs(
    image_uri=image_uri,
    args=args,
    replica_count=replica_count,
    machine_type=machine_type,
    accelerator_count=accelerator_count,
    accelerator_type=accelerator_type
)

pp = pprint.PrettyPrinter()
print(pp.pformat(worker_pool_specs))

# +
display_name = 'custom-test'

job = aiplatform.CustomJob(
    display_name=display_name,
    worker_pool_specs=worker_pool_specs,
)

job.run(sync=False)
# -

job.resource_name

job.wait()



StarcoderdataJupyterScriptsDedupFiltered