Added dask_bring_your_own_container_local_processing Sample Code

Author: eitansela

dask_bring_your_own_container_local_processing/container/Dockerfile

@@ -0,0 +1,23 @@
+FROM continuumio/miniconda3:4.7.12
+
+ENV PYTHONHASHSEED 0
+ENV PYTHONIOENCODING UTF-8
+
+# Install required Python packages fo Dask
+RUN conda install --yes dask distributed dask-ml boto3
+
+# Install additional Python packages
+RUN conda install aiohttp boto3
+
+# Dumb init
+RUN wget -O /usr/local/bin/dumb-init https://github.com/Yelp/dumb-init/releases/download/v1.2.0/dumb-init_1.2.0_amd64
+RUN chmod +x /usr/local/bin/dumb-init
+
+RUN mkdir /opt/app /etc/dask
+COPY dask_config/dask.yaml /etc/dask/
+
+# Set up bootstrapping program and Dask configuration
+COPY program /opt/program
+RUN chmod +x /opt/program/bootstrap.py
+
+ENTRYPOINT ["/opt/program/bootstrap.py"]

dask_bring_your_own_container_local_processing/container/dask_config/dask.yaml

@@ -0,0 +1,11 @@
+logging:
+ distributed: info
+ bokeh: critical
+ tornado: critical
+
+scheduler:
+ work-stealing: True
+ allowed-failures: 10
+
+admin:
+ log-format: '%(name)s - %(levelname)s - %(message)s'

dask_bring_your_own_container_local_processing/container/program/bootstrap.py

@@ -0,0 +1,94 @@
+#!/usr/bin/env python3
+import json
+import os
+import socket
+import subprocess
+import sys
+import time
+from shutil import copyfile
+from subprocess import PIPE, Popen
+
+DASK_PATH = "/opt/conda/bin"
+
+
+def get_resource_config():
+ resource_config_path = "/opt/ml/config/resourceconfig.json"
+ with open(resource_config_path, "r") as f:
+ resource_config_json = json.load(f)
+ return resource_config_json
+
+
+def start_daemons(master_ip):
+ resource_config = get_resource_config()
+ print(f"resource_config: {resource_config}")
+ current_host = resource_config["current_host"]
+ print(f"current_host: {current_host}")
+ scheduler_host = resource_config["hosts"][0]
+ print(f"scheduler_host: {scheduler_host}")
+
+
+ cmd_start_scheduler = os.path.join(DASK_PATH, "dask-scheduler")
+ cmd_start_worker = os.path.join(DASK_PATH, "dask-worker")
+ schedule_conn_string = "tcp://{ip}:8786".format(ip=master_ip)
+ if current_host == scheduler_host:
+ Popen([cmd_start_scheduler])
+ Popen([cmd_start_worker, schedule_conn_string])
+ else:
+ worker_process = Popen([cmd_start_worker, schedule_conn_string])
+
+
+def get_ip_from_host(host_name):
+ ip_wait_time = 200
+ counter = 0
+ ip = ""
+
+ while counter < ip_wait_time and ip == "":
+ try:
+ ip = socket.gethostbyname(host_name)
+ break
+ except:
+ counter += 1
+ time.sleep(1)
+
+ if counter == ip_wait_time and ip == "":
+ raise Exception(
+ "Exceeded max wait time of {}s for hostname resolution".format(ip_wait_time)
+ )
+
+ return ip
+
+
+if __name__ == "__main__":
+ ips = []
+ resource_config = get_resource_config()
+ master_host = resource_config["hosts"][0]
+ scheduler_ip = get_ip_from_host(master_host)
+ current_host = resource_config["current_host"]
+
+ # Start Dask cluster in all nodes
+ start_daemons(scheduler_ip)
+
+ # wait for a few seconds until the cluster is up and running
+ time.sleep(10)
+
+ # Submit the preprocessing job on the cluster from the first instance. You only need to submit the job once from one node.
+ if current_host == master_host:
+ cmd_string = ["/opt/conda/bin/python", str(sys.argv[1])]
+ cmd_string.extend(sys.argv[2:])
+ cmd_string.append(scheduler_ip)
+ print(f"cmd_string: {cmd_string}")
+ result = subprocess.Popen(cmd_string)
+ _ = result.communicate()[0]
+ exit_code = result.returncode
+ else:
+ while True:
+ scheduler = (scheduler_ip, 8786)
+ alive_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+ alive_check = alive_socket.connect_ex(scheduler)
+ if alive_check == 0:
+ pass
+ else:
+ print("Received a shutdown signal from Dask cluster")
+ sys.exit(0)
+ alive_socket.close()
+ time.sleep(2)

dask_bring_your_own_container_local_processing/dask_bring_your_own_container_local_processing.py

@@ -0,0 +1,54 @@
+# This is a sample Python program that runs a Dask Processing job on a JSON fetched from a web site.
+# The output of the processing will be total files found in the JSON.
+# This implementation will work on your *local computer*.
+#
+# Prerequisites:
+# 1. Install required Python packages:
+# pip install boto3 sagemaker pandas scikit-learn
+# pip install 'sagemaker[local]'
+# 2. Docker Desktop has to be installed on your computer, and running.
+# 3. Open terminal and run the following commands:
+# docker build -t sagemaker-dask-processing-local container/.
+########################################################################################################################
+
+from sagemaker.local import LocalSession
+from sagemaker.processing import ScriptProcessor, ProcessingInput, ProcessingOutput
+import boto3
+
+
+s3 = boto3.client('s3')
+sagemaker_session = LocalSession()
+sagemaker_session.config = {'local': {'local_code': True}}
+
+# For local training a dummy role will be sufficient
+role = 'arn:aws:iam::111111111111:role/service-role/AmazonSageMaker-ExecutionRole-20200101T000001'
+
+dask_processor = ScriptProcessor(command=["/opt/program/bootstrap.py"],
+ image_uri='sagemaker-dask-processing-local',
+ role=role,
+ instance_count=1,
+ instance_type='local')
+
+dask_processor.run(code='processing_script.py',
+ outputs=[ProcessingOutput(
+ output_name='filenames_processed_data',
+ source='/opt/ml/processing/processed_data/')],
+ arguments=['site_uri', 'https://archive.analytics.mybinder.org/index.jsonl']
+ )
+
+preprocessing_job_description = dask_processor.jobs[-1].describe()
+output_config = preprocessing_job_description['ProcessingOutputConfig']
+
+print(output_config)
+
+for output in output_config['Outputs']:
+ if output['OutputName'] == 'filenames_processed_data':
+ filenames_processed_data_file = output['S3Output']['S3Uri']
+ bucket = filenames_processed_data_file.split("/")[:3][2]
+ output_file_name = '/'.join(filenames_processed_data_file.split("/")[3:])+"/filenames_in_json.txt"
+
+print(f'Opening processing output file: {"s3://"+bucket+"/"+output_file_name}')
+data = s3.get_object(Bucket=bucket, Key=output_file_name)
+contents = data['Body'].read()
+print('Processing output file content\n-----------\n')
+print(contents.decode("utf-8"))

dask_bring_your_own_container_local_processing/processing_script.py

@@ -0,0 +1,55 @@
+import os
+import sys
+import logging
+from dask.distributed import Client
+import dask.bag as db
+import json
+
+
+processed_data_path = '/opt/ml/processing/processed_data'
+
+
+def main():
+ print("Processing Started")
+
+ # Convert command line args into a map of args
+ args_iter = iter(sys.argv[1:])
+ args = dict(zip(args_iter, args_iter))
+ scheduler_ip = sys.argv[-1]
+ print(f"scheduler_ip: {scheduler_ip}")
+
+ # Start the Dask cluster client
+ try:
+ print("initiating client")
+ client = Client("tcp://{ip}:8786".format(ip=scheduler_ip))
+ print("Cluster information: {}".format(client))
+ except Exception as err:
+ logging.exception(err)
+
+ print(f"Received arguments {args}")
+
+ if args["site_uri"]:
+ print(f"Processing web site JSON: {args['site_uri']}")
+ filenames = (db.read_text(args['site_uri'])
+ .map(json.loads)
+ .pluck('name')
+ .compute())
+
+ filenames = ['https://archive.analytics.mybinder.org/' + fn for fn in filenames]
+ print(f"Total filenames: {len(filenames)}")
+ print(f"Sample filenames found: {filenames[:5]}")
+
+ output_file = os.path.join(processed_data_path, "filenames_in_json.txt")
+ print(f'Writing output file: {output_file}')
+ with open(output_file, 'w') as outfile:
+ outfile.write(json.dumps(filenames))
+ else:
+ print("No `site_uri` parameter - doing nothing")
+
+ print("Processing Complete")
+
+ print(client)
+ sys.exit(os.EX_OK)
+
+if __name__ == "__main__":
+ main()