Step Functionsの動的並列処理をローカルで実行する

import os
import boto3
from scatter import ScatterHandler

if os.getenv("LOCALSTACK_HOSTNAME") is None:
    s3 = boto3.client("s3", "ap-northeast-1")
else:
    endpoint = f"http://{os.environ['LOCALSTACK_HOSTNAME']}:4566"
    s3 = boto3.client(
        service_name="s3",
        endpoint_url=endpoint,
        aws_access_key_id="test",
        aws_secret_access_key="test",
    )


def lambda_handler(event, context) -> dict:
    handler = ScatterHandler(
        event,
        context,
        s3,
    )
    return handler.main()

import json
import tempfile
import logging
import pandas as pd
from typing import List, Dict

logger = logging.getLogger()
logger.setLevel(logging.INFO)


class ScatterHandler(object):
    def __init__(
        self,
        event,
        context,
        s3,
    ):
        self.event = event
        self.context = context
        self.s3 = s3

    def main(self) -> dict:
        try:
            bucket = "test-bucket"
            data_path = self.event["input_obj"]
            division_number = 10
            segments: Dict = {}
            segments["segment_definitions"] = []
            data = self.get_s3_data(bucket, data_path)
            df = self.make_df(data)
            dfs = [
                df.loc[i : i + division_number - 1, :]
                for i in range(0, len(df), division_number)
            ]
            segments = self.make_segment_df(segments, bucket, dfs)
            return segments

        except Exception as e:
            logger.exception(e)
            raise e

    def get_s3_data(self, bucket: str, key: str) -> List[dict]:
        resp = self.s3.get_object(Bucket=bucket, Key=key)
        body = resp["Body"].read().decode("utf-8")
        json_dict: List[dict] = json.loads(body)
        return json_dict

    def make_df(self, data: list) -> pd.DataFrame:
        df = pd.DataFrame.from_dict(data)
        return df

    def make_segment_df(self, segments: dict, bucket: str, dfs: list) -> dict:
        for i, df_i in enumerate(dfs):
            with tempfile.TemporaryFile() as fp:
                df_i.to_pickle(fp)
                fp.seek(0)
                fsend = "scatter/job_" + str(i).zfill(3) + ".pkl"
                self.s3.put_object(
                    Body=fp.read(),
                    Bucket=bucket,
                    Key=fsend,
                )
                segments["segment_definitions"].append(fsend)
        return segments

import tempfile
import logging
import pickle
import pandas as pd

logger = logging.getLogger()
logger.setLevel(logging.INFO)


class SegmentHandler(object):
    def __init__(self, event, context, s3):
        self.event = event
        self.context = context
        self.s3 = s3

    def main(
        self,
    ) -> str:
        try:
            bucket = "test-bucket"
            recieve = self.event
            send = recieve.replace("scatter", "gather")
            df = self.get_s3_data(bucket, recieve)
            df = self.process(df)
            return self.send_segment_df(df, bucket, send)

        except Exception as e:
            logger.exception(e)
            raise e

    def get_s3_data(self, bucket, key) -> pd.DataFrame:
        resp = self.s3.get_object(Bucket=bucket, Key=key)
        body = resp["Body"].read()
        df = pickle.loads(body)
        return df

    def calc(self, row):
        if row["会員ランク"] > 3:
            return row["ポイント"] * 1.25
        else:
            return row["ポイント"]

    def process(self, data: pd.DataFrame) -> pd.DataFrame:
        data["ボーナスポイント"] = data.apply(self.calc, axis=1)
        return data

    def make_df(self, data: list) -> pd.DataFrame:
        df = pd.DataFrame.from_dict(data)
        return df

    def send_segment_df(self, df: pd.DataFrame, bucket: str, send: str) -> str:
        with tempfile.TemporaryFile() as fp:
            df.to_pickle(fp)
            fp.seek(0)
            self.s3.put_object(
                Body=fp.read(),
                Bucket=bucket,
                Key=send,
            )
        return send

import os
import boto3
from segment import SegmentHandler

if os.getenv("LOCALSTACK_HOSTNAME") is None:
    s3 = boto3.client("s3", "ap-northeast-1")
else:
    endpoint = f"http://{os.environ['LOCALSTACK_HOSTNAME']}:4566"
    s3 = boto3.client(
        service_name="s3",
        endpoint_url=endpoint,
        aws_access_key_id="test",
        aws_secret_access_key="test",
    )


def lambda_handler(event, context) -> str:
    handler = SegmentHandler(event, context, s3)
    return handler.main()

import tempfile
import logging
import pandas as pd
import pickle
from typing import List

logger = logging.getLogger()
logger.setLevel(logging.INFO)


class GatherHandler(object):
    def __init__(self, event, context, s3):
        self.event = event
        self.context = context
        self.s3 = s3

    def main(self) -> str:
        try:
            bucket = "test-bucket"
            segments = self.event["segment_results"]
            send = "test.xlsx"
            data_frames: List[pd.DataFrame] = []
            for pkl in segments:
                df = self.get_s3_df(bucket, pkl)
                data_frames.append(df)
            df_gather = pd.concat(data_frames)
            return self.send_excel(df_gather, bucket, send)

        except Exception as e:
            logger.exception(e)
            raise e

    def get_s3_df(self, bucket, key) -> pd.DataFrame:
        resp = self.s3.get_object(Bucket=bucket, Key=key)
        body = resp["Body"].read()
        df = pickle.loads(body)
        return df

    def send_excel(self, df: pd.DataFrame, bucket: str, send: str) -> str:
        with tempfile.TemporaryFile() as fp:
            writer = pd.ExcelWriter(fp, engine="xlsxwriter")
            df.to_excel(writer, sheet_name="Sheet1", index=False)
            writer.save()
            fp.seek(0)
            self.s3.put_object(
                Body=fp.read(),
                Bucket=bucket,
                Key=send,
            )
        return send

import os
import boto3
from gather import GatherHandler

if os.getenv("LOCALSTACK_HOSTNAME") is None:
    s3 = boto3.client("s3", "ap-northeast-1")
else:
    endpoint = f"http://{os.environ['LOCALSTACK_HOSTNAME']}:4566"
    s3 = boto3.client(
        service_name="s3",
        endpoint_url=endpoint,
        aws_access_key_id="test",
        aws_secret_access_key="test",
    )


def lambda_handler(event, context) -> str:
    handler = GatherHandler(event, context, s3)
    return handler.main()

import datetime
import json
from random import randint

import boto3
from fire import Fire
from mimesis import Person
from mimesis.locales import Locale

person = Person(Locale.JA)


def dummy_data(num: int) -> dict:
    id = str(num).zfill(3)
    date = datetime.date(2021, randint(1, 12), randint(1, 28)).strftime("%Y-%m-%d")
    dummy_dict = {
        "会員番号": id,
        "名前": person.full_name(reverse=True),
        "会員ランク": randint(1, 5),
        "ポイント": randint(50, 100),
        "タイムスタンプ": date,
    }
    return dummy_dict


def send_json(s3, sample_data: list, bucket: str, send: str) -> str:
    with open("utils/data/sample_data.json", mode="wt", encoding="utf-8") as f:
        json.dump(sample_data, f, ensure_ascii=False, indent=4)
    s3.put_object(
        Body=json.dumps(sample_data, ensure_ascii=False, indent=4),
        Bucket=bucket,
        Key=send,
    )
    return send


def make_dummy_data(k) -> list:
    sample_data = []
    for i in range(k):
        sample_data.append(dummy_data(i))
    return sample_data


def main(iterate_num: int) -> str:
    endpoint = f"http://localhost:4566"
    s3 = boto3.client(
        service_name="s3",
        endpoint_url=endpoint,
        aws_access_key_id="test",
        aws_secret_access_key="test",
    )
    bucket = "test-bucket"
    send = "test.json"
    sample_data = make_dummy_data(iterate_num)
    send = send_json(s3, sample_data, bucket, send)
    return send


if __name__ == "__main__":
    Fire(main)