#

# Licensed to the Apache Software Foundation (ASF) under one or more

# contributor license agreements. See the NOTICE file distributed with

# this work for additional information regarding copyright ownership.

# The ASF licenses this file to You under the Apache License, Version 2.0

# (the "License"); you may not use this file except in compliance with

# the License. You may obtain a copy of the License at

#

# http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

#

import datetime

from operator import itemgetter

import joblib

import numpy as np

import pandas as pd

from sklearn.cluster import MiniBatchKMeans

import apache_beam as beam

from apache_beam.coders import PickleCoder

from apache_beam.coders import VarIntCoder

from apache_beam.io.filesystems import FileSystems

from apache_beam.ml.inference.base import PredictionResult

from apache_beam.ml.inference.base import RunInference

from apache_beam.ml.inference.sklearn_inference import ModelFileType

from apache_beam.ml.inference.sklearn_inference import SklearnModelHandlerNumpy

from apache_beam.transforms import core

from apache_beam.transforms import ptransform

from apache_beam.transforms.userstate import ReadModifyWriteStateSpec

class SaveModel(core.DoFn):

"""Saves trained clustering model to persistent storage"""

def __init__(self, checkpoints_path: str):

self.checkpoints_path = checkpoints_path

def process(self, model):

# generate ISO 8601

iso_timestamp = datetime.datetime.now(

datetime.timezone.utc).strftime("%Y%m%dT%H%M%SZ")

checkpoint_name = f'{self.checkpoints_path}/{iso_timestamp}.checkpoint'

latest_checkpoint = f'{self.checkpoints_path}/latest.checkpoint'

# rename previous checkpoint

if FileSystems.exists(latest_checkpoint):

FileSystems.rename([latest_checkpoint], [checkpoint_name])

file = FileSystems.create(latest_checkpoint, 'wb')

if not joblib:

raise ImportError(

'Could not import joblib in this execution environment. '

'For help with managing dependencies on Python workers.'

'see https://beam.apache.org/documentation/sdks/python-pipeline-dependencies/' # pylint: disable=line-too-long

)

joblib.dump(model, file)

yield checkpoint_name

class AssignClusterLabelsFn(core.DoFn):

"""Takes a trained model and input data and labels

all data instances using the trained model."""

def process(self, batch, model, model_id):

cluster_labels = model.predict(batch)

for e, i in zip(batch, cluster_labels):

yield PredictionResult(example=e, inference=i, model_id=model_id)

class SelectLatestModelState(core.CombineFn):

"""Selects that latest version of a model after training"""

def create_accumulator(self):

# create and initialise accumulator

return None, 0

def add_input(self, accumulator, element):

# accumulates each element from input in accumulator

if element[1] > accumulator[1]:

return element

return accumulator

def merge_accumulators(self, accumulators):

# Multiple accumulators could be processed in parallel,

# this function merges them

return max(accumulators, key=itemgetter(1))

def extract_output(self, accumulator):

# Only output the tracker

return accumulator[0]

class ClusteringAlgorithm(core.DoFn):

"""Abstract class with the interface

that clustering algorithms need to follow."""

MODEL_SPEC = ReadModifyWriteStateSpec("clustering_model", PickleCoder())

ITERATION_SPEC = ReadModifyWriteStateSpec(

'training_iterations', VarIntCoder())

MODEL_ID = 'ClusteringAlgorithm'

def __init__(

self, n_clusters: int, checkpoints_path: str, cluster_args: dict):

super().__init__()

self.n_clusters = n_clusters

self.checkpoints_path = checkpoints_path

self.cluster_args = cluster_args

self.clustering_algorithm = None

def process(

self,

keyed_batch,

model_state=core.DoFn.StateParam(MODEL_SPEC),

iteration_state=core.DoFn.StateParam(ITERATION_SPEC),

*args,

**kwargs):

raise NotImplementedError

def load_model_checkpoint(self):

latest_checkpoint = f'{self.checkpoints_path}/latest.checkpoint'

if FileSystems.exists(latest_checkpoint):

file = FileSystems.open(latest_checkpoint, 'rb')

if not joblib:

raise ImportError(

'Could not import joblib in this execution environment. '

'For help with managing dependencies on Python workers.'

'see https://beam.apache.org/documentation/sdks/python-pipeline-dependencies/' # pylint: disable=line-too-long

)

return joblib.load(file)

return self.clustering_algorithm(

n_clusters=self.n_clusters, **self.cluster_args)

class OnlineKMeans(ClusteringAlgorithm):

"""Online K-Means function. Used the MiniBatchKMeans from sklearn

More information: https://scikit-learn.org/stable/modules/generated/sklearn.cluster.MiniBatchKMeans.html""" # pylint: disable=line-too-long

MODEL_SPEC = ReadModifyWriteStateSpec("clustering_model", PickleCoder())

ITERATION_SPEC = ReadModifyWriteStateSpec(

'training_iterations', VarIntCoder())

MODEL_ID = 'OnlineKmeans'

def __init__(

self, n_clusters: int, checkpoints_path: str, cluster_args: dict):

super().__init__(n_clusters, checkpoints_path, cluster_args)

self.clustering_algorithm = MiniBatchKMeans

def process(

self,

keyed_batch,

model_state=core.DoFn.StateParam(MODEL_SPEC),

iteration_state=core.DoFn.StateParam(ITERATION_SPEC),

*args,

**kwargs):

# 1. Initialise or load states

clustering = model_state.read() or self.load_model_checkpoint()

iteration = iteration_state.read() or 0

iteration += 1

# 2. Remove the temporary assigned keys

_, batch = keyed_batch

# 3. Calculate cluster centroids

clustering.partial_fit(batch)

# 4. Store the training set and model

model_state.write(clustering)

iteration_state.write(iteration)

# checkpoint = joblib.dump(clustering, f'kmeans_checkpoint_{iteration}')

yield clustering, iteration

class ConvertToNumpyArray(core.DoFn):

"""Helper function to convert incoming data

to numpy arrays that are accepted by sklearn"""

def process(self, element, *args, **kwargs):

if isinstance(element, (tuple, list)):

yield np.array(element)

elif isinstance(element, np.ndarray):

yield element

elif isinstance(element, (pd.DataFrame, pd.Series)):

yield element.to_numpy()

else:

raise ValueError(f"Unsupported type: {type(element)}")

class ClusteringPreprocessing(ptransform.PTransform):

def __init__(

self, n_clusters: int, batch_size: int, is_batched: bool = False):

""" Preprocessing for Clustering Transformation

The clustering transform expects batches for performance reasons,

therefore this batches the data and converts it to numpy arrays,

which are accepted by sklearn. This transform also adds the same key

to all batches, such that only 1 state is created and updated during

clustering updates.

Example Usage::

pcoll | ClusteringPreprocessing(

n_clusters=8,

batch_size=1024,

is_batched=False)

Args:

n_clusters: number of clusters used by the algorithm

batch_size: size of the data batches

is_batched: boolean value that marks if the collection is already

batched and thus doesn't need to be batched by this transform

"""

super().__init__()

self.n_clusters = n_clusters

self.batch_size = batch_size

self.is_batched = is_batched

def expand(self, pcoll):

pcoll = (

pcoll

|

"Convert element to numpy arrays" >> beam.ParDo(ConvertToNumpyArray()))

if not self.is_batched:

pcoll = (

pcoll

| "Create batches of elements" >> beam.BatchElements(

min_batch_size=self.n_clusters, max_batch_size=self.batch_size)

| "Covert to 2d numpy array" >>

beam.Map(lambda record: np.array(record)))

return pcoll

class OnlineClustering(ptransform.PTransform):

def __init__(

self,

clustering_algorithm,

n_clusters: int,

cluster_args: dict,

checkpoints_path: str,

batch_size: int = 1024,

is_batched: bool = False):

""" Clustering transformation itself, it first preprocesses the data,

then it applies the clustering transformation step by step on each

of the batches.

Example Usage::

pcoll | OnlineClustering(

clustering_algorithm=OnlineKMeansClustering

batch_size=1024,

n_clusters=6

cluster_args={}))

Args:

clustering_algorithm: Clustering algorithm (DoFn)

n_clusters: Number of clusters

cluster_args: Arguments for the sklearn clustering algorithm

(check sklearn documentation for more information)

batch_size: size of the data batches

is_batched: boolean value that marks if the collection is already

batched and thus doesn't need to be batched by this transform

"""

super().__init__()

self.clustering_algorithm = clustering_algorithm

self.n_clusters = n_clusters

self.batch_size = batch_size

self.cluster_args = cluster_args

self.checkpoints_path = checkpoints_path

self.is_batched = is_batched

def expand(self, pcoll):

# 1. Preprocess data for more efficient clustering

data = (

pcoll

| 'Batch data for faster processing' >> ClusteringPreprocessing(

n_clusters=self.n_clusters,

batch_size=self.batch_size,

is_batched=self.is_batched)

| "Add a key for stateful processing" >>

beam.Map(lambda record: (1, record)))

# 2. Calculate cluster centers

model = (

data

| 'Cluster' >> core.ParDo(

self.clustering_algorithm(

n_clusters=self.n_clusters,

cluster_args=self.cluster_args,

checkpoints_path=self.checkpoints_path))

| 'Select latest model state' >> core.CombineGlobally(

SelectLatestModelState()).without_defaults())

# 3. Save the trained model checkpoint to persistent storage,

# so it can be loaded for further training in the next window

# or loaded into an sklearn modelhandler for inference

_ = (model | core.ParDo(SaveModel(checkpoints_path=self.checkpoints_path)))

return model

class AssignClusterLabelsRunInference(ptransform.PTransform):

def __init__(self, checkpoints_path):

super().__init__()

self.clustering_model = SklearnModelHandlerNumpy(

model_uri=f'{checkpoints_path}/latest.checkpoint',

model_file_type=ModelFileType.JOBLIB)

def expand(self, pcoll):

predictions = (

pcoll

| "RunInference" >> RunInference(self.clustering_model))

return predictions

class AssignClusterLabelsInMemoryModel(ptransform.PTransform):

def __init__(

self, model, n_clusters, batch_size, is_batched=False, model_id=None):

self.model = model

self.n_clusters = n_clusters

self.batch_size = batch_size

self.is_batched = is_batched

self.model_id = model_id

def expand(self, pcoll):

return (

pcoll

| "Preprocess data for faster prediction" >> ClusteringPreprocessing(

n_clusters=self.n_clusters,

batch_size=self.batch_size,

is_batched=self.is_batched)

| "Assign cluster labels" >> core.ParDo(

AssignClusterLabelsFn(), model=self.model, model_id=self.model_id))