Spark MLlib on AWS Glue. Distributed ML on AWS that’s ready to... | by John Elliott | Towards Data Science

AWS pushes Sagemaker as its machine learning platform. However, Spark’s MLlib is a comprehensive library that runs distributed ML natively on AWS Glue — and provides a viable alternative to their primary ML platform.
One of the big benefits of Sagemaker is that it easily supports experimentation via its Jupyter Notebooks. But operationalising your Sagemaker ML can be difficult, particularly if you need to include ETL processing at the start of your pipeline. In this situation, Apache Spark’s MLlib running on AWS Glue can be a good option — by its very nature, it is immediately operationalised, integrated with ETL pre-processing and ready to be used in production for an end-to-end machine learning pipeline.
By its very nature, [AWS Glue] is immediately operationalised, integrated with ETL pre-processing and ready to be used in production for an end-to-end machine learning pipeline
AWS Glue is a managed Spark ETL platform for processing large volumes of data via distributed machines. MLlib comes as part of Spark 2.4, which is the default version on AWS Glue. There is no need to add libraries to use MLlib within AWS Glue.
As a distributed ETL platform, AWS Glue (via Spark) allows you to perform your data pre-processing at large scale easily. Glue Studio provides a nice UI for building directed acyclic graphs that represent the flow of data through each pre-processing step (see above image for an example). From the graph UI, you can rename fields, convert types, filter records and remove columns. However, some of the more complex data processing activities require a ‘Custom Transform’ to be added. Actions that require a Custom Transform include:
Running MLlib functions, such as TF-IDF or models such as Random Forest or Gradient Boosted Tree algorithmsAggregating values, such as generating count, mean and sum valuesComplex column transforms such as one-hot encoding or concatenating values
Running MLlib functions, such as TF-IDF or models such as Random Forest or Gradient Boosted Tree algorithms
Aggregating values, such as generating count, mean and sum values
Complex column transforms such as one-hot encoding or concatenating values
After performing your pre-processing and featurization using ETL transforms, you need to train your data. To access MLlib in AWS Glue, you just need to add the import statements and you’re ready to go. Below is an example of how to train a regression Decision Tree model in a AWS Glue Studio Custom Transform using PySpark:
# Get the dataframe. Ensure there's a 'features' column.df = dfc.select(list(dfc.keys())[0]).toDF()# Get the logger for Cloudwatch Logslogger = glueContext.get_logger()from pyspark.ml import Pipelinefrom pyspark.ml.regression import DecisionTreeRegressorfrom pyspark.ml.evaluation import RegressionEvaluator# Split data into training and test sets(trainingData, testData) = df.randomSplit([0.7, 0.3])# Create a DecisionTree model.dt = DecisionTreeRegressor(featuresCol="features")# Create a pipeline to wrap the DecisionTree pipeline = Pipeline(stages=[dt])# Train model.model = pipeline.fit(trainingData)# Make predictions.predictions = model.transform(testData)# Select (prediction, true label) and compute test errorevaluator = RegressionEvaluator(    labelCol="label", predictionCol="prediction", metricName="rmse")rmse = evaluator.evaluate(predictions)logger.info("Root Mean Squared Error (RMSE) on test data = %g" % rmse)
In the above code, the root-mean-squared error (RMSE) is logged out to Cloudwatch Logs. To use the model to predict, the simplest option is to use the same Custom Transform and run model.transform on a dataframe of unseen inference data. To then use the predictions, you will need to return them from the Custom Transform in a DynamicFrameCollection (see code snippet at the bottom of this article).
There are a few gotchas to look out for when running Spark’s MLlib in AWS Glue:
Experimentation can be slow and difficult. Iterating and debugging can be frustrating because at a minimum each job takes one minute to run. When experimenting, run on a reduced data set to speed up the process (it will still take 1 minute+). Reduce the number of workers to 2 and use G.1X to save money. Disable bookmarks and retries.To help with debugging, you can log out information to Cloudwatch from your Custom Transform. You will have to go to the /aws-glue/jobs/logs-v2 log group on Cloudwatch, then open the log stream that ends with ‘-driver’ to see the logged-out values. Below is a PySpark example of a Glue Studio Custom Transform with Cloudwatch logging set up.AWS Glue Developer Endpoints may help with experimentation and debugging. These allow you to run a development notebook on a Spark cluster and speed up your development iterations when writing code. The main downside is that they are expensive (and can take 5–15 minutes to spin-up the cluster). Note that they charge by the second, so can be significantly more expensive than just executing your jobs directly. One option is to use Developer Endpoints while you are learning Spark, then move over to directly executing jobs in AWS Glue after you are familiar with writing Spark code.While you can run your ML training and predictions in the same job as your featurization ETL, you may want to separate them into two jobs. By separating them into two jobs, you can separate their configurations. AWS Glue provides two worker types: G.1X and G.2X. You may find that G.1X is adequate for ETL, but G.2X may be more suitable for ML training and inference as each worker has twice as much memory available.Only batch inference works on AWS Glue — real-time inference isn’t an option here. Use Sagemaker endpoints for that use case.MLlib on AWS Glue easily allows you to retrain your whole dataset, then perform inference in a single job. Training is probably the most expensive step, however, you can benefit from better predictions by retraining before batch inferenceTry to use native Spark dataframes and libraries, rather than, for example, Pandas dataframes. This will allow you to leverage the distributed capabilities of Spark much better.There are two APIs available in MLlib — the dataframe-based API and the RDD-based API. You should use the dataframe based approaches as much as possible. Likewise, try to avoid using user-defined functions (UDFs). Both UDFs and RDD-based capabilities are slower than native dataframe methods.SparseVector is an efficient way of representing sparse datasets (i.e. data sources with a large number of columns and most values being zero). A common example of this is the output of TF-IDF when performing natural language processing. Using a SparseVector column avoids adding hundreds or thousands of columns to a dataframe (which would most likely cause your job to fail due to out-of-memory or timeout).Some dataframe column types created by MLlib cannot be written out of a Custom Transform without being cast to integer or string first. You may get an error such as ‘An error occurred while calling o341.toDF. org.apache.spark.ml.linalg.VectorUDT@3bfc3ba7 (of class org.apache.spark.ml.linalg.VectorUDT)’ when this happens.In AWS Glue reading to and from S3 is much easier than any other data source / destination. Directly connecting to a database may be more trouble than it’s worth.You will most likely need to make use of Glue Crawlers to learn your schema on S3. You will need to re-run the crawler each time your data changes. Use partition keys in your S3 path/key and partition predicates on your AWS Glue ‘S3 data source’ input node to filter the data to a specific date.JSON and CSV can be useful data types when developing as you can see the inputs and outputs. However, they are slow and large. For production, consider switching to Parquet format inputs and outputs.Exporting a SparseVector column (that has been cast to String) to CSV leads to an invalid CSV structure. Parquet format works instead.Be careful to distinguish between the new dataframe-based API and the older RDD-based API. Ensure you import the correct classes. In PySpark, you should import ‘pyspark.ml.linalg’ rather than ‘pyspark.mllib.linalg’ to access the newer API. If you mix-up types, you may get strange errors such as “IllegalArgumentException: ‘requirement failed: Column features must be of type struct<type:tinyint,size:int,indices:array<int>,values:array<double>> but was actually struct<type:tinyint,size:int,indices:array<int>,values:array<double>>.’. Both of those values look the same but the difference is that they’re using an incompatible version of the object from the wrong API.
Experimentation can be slow and difficult. Iterating and debugging can be frustrating because at a minimum each job takes one minute to run. When experimenting, run on a reduced data set to speed up the process (it will still take 1 minute+). Reduce the number of workers to 2 and use G.1X to save money. Disable bookmarks and retries.
To help with debugging, you can log out information to Cloudwatch from your Custom Transform. You will have to go to the /aws-glue/jobs/logs-v2 log group on Cloudwatch, then open the log stream that ends with ‘-driver’ to see the logged-out values. Below is a PySpark example of a Glue Studio Custom Transform with Cloudwatch logging set up.
AWS Glue Developer Endpoints may help with experimentation and debugging. These allow you to run a development notebook on a Spark cluster and speed up your development iterations when writing code. The main downside is that they are expensive (and can take 5–15 minutes to spin-up the cluster). Note that they charge by the second, so can be significantly more expensive than just executing your jobs directly. One option is to use Developer Endpoints while you are learning Spark, then move over to directly executing jobs in AWS Glue after you are familiar with writing Spark code.
While you can run your ML training and predictions in the same job as your featurization ETL, you may want to separate them into two jobs. By separating them into two jobs, you can separate their configurations. AWS Glue provides two worker types: G.1X and G.2X. You may find that G.1X is adequate for ETL, but G.2X may be more suitable for ML training and inference as each worker has twice as much memory available.
Only batch inference works on AWS Glue — real-time inference isn’t an option here. Use Sagemaker endpoints for that use case.
MLlib on AWS Glue easily allows you to retrain your whole dataset, then perform inference in a single job. Training is probably the most expensive step, however, you can benefit from better predictions by retraining before batch inference
Try to use native Spark dataframes and libraries, rather than, for example, Pandas dataframes. This will allow you to leverage the distributed capabilities of Spark much better.
There are two APIs available in MLlib — the dataframe-based API and the RDD-based API. You should use the dataframe based approaches as much as possible. Likewise, try to avoid using user-defined functions (UDFs). Both UDFs and RDD-based capabilities are slower than native dataframe methods.
SparseVector is an efficient way of representing sparse datasets (i.e. data sources with a large number of columns and most values being zero). A common example of this is the output of TF-IDF when performing natural language processing. Using a SparseVector column avoids adding hundreds or thousands of columns to a dataframe (which would most likely cause your job to fail due to out-of-memory or timeout).
Some dataframe column types created by MLlib cannot be written out of a Custom Transform without being cast to integer or string first. You may get an error such as ‘An error occurred while calling o341.toDF. org.apache.spark.ml.linalg.VectorUDT@3bfc3ba7 (of class org.apache.spark.ml.linalg.VectorUDT)’ when this happens.
In AWS Glue reading to and from S3 is much easier than any other data source / destination. Directly connecting to a database may be more trouble than it’s worth.
You will most likely need to make use of Glue Crawlers to learn your schema on S3. You will need to re-run the crawler each time your data changes. Use partition keys in your S3 path/key and partition predicates on your AWS Glue ‘S3 data source’ input node to filter the data to a specific date.
JSON and CSV can be useful data types when developing as you can see the inputs and outputs. However, they are slow and large. For production, consider switching to Parquet format inputs and outputs.
Exporting a SparseVector column (that has been cast to String) to CSV leads to an invalid CSV structure. Parquet format works instead.
Be careful to distinguish between the new dataframe-based API and the older RDD-based API. Ensure you import the correct classes. In PySpark, you should import ‘pyspark.ml.linalg’ rather than ‘pyspark.mllib.linalg’ to access the newer API. If you mix-up types, you may get strange errors such as “IllegalArgumentException: ‘requirement failed: Column features must be of type struct<type:tinyint,size:int,indices:array<int>,values:array<double>> but was actually struct<type:tinyint,size:int,indices:array<int>,values:array<double>>.’. Both of those values look the same but the difference is that they’re using an incompatible version of the object from the wrong API.
Yes, Spark ETL via AWS Glue can be integrated with Amazon Sagemaker. A typical workflow might be:
Experiment and train the model in Sagemaker using Jupyter NotebooksProductionize the model by deploying a batch inference model from Sagemaker notebooksProductionize pre-processing and featurization using ETL in AWS GlueSchedule / trigger a batch inference model to run after completion of the AWS Glue ETL
Experiment and train the model in Sagemaker using Jupyter Notebooks
Productionize the model by deploying a batch inference model from Sagemaker notebooks
Productionize pre-processing and featurization using ETL in AWS Glue
Schedule / trigger a batch inference model to run after completion of the AWS Glue ETL
The above approach may seem to be optimal use of AWS services for their intended purposes. However, you do need to look out for some things before going down this path. Some column types, for example, are not compatible between Spark MLlib and Sagemaker models. For instance, Sagemaker XGBoost requires its inputs to be in a specific format and cannot read SparseVector columns from Spark.
Below is an example of a PySpark Custom Transform on AWS Glue Studio for logging out to the ‘-driver’ log stream under the /aws-glue/jobs/logs-v2 log stream.
def MyTransform (glueContext, dfc) -> DynamicFrameCollection:     logger = glueContext.get_logger()     df = dfc.select(list(dfc.keys())[0]).toDF()     logger.info("Number of df rows:" + str(df.count()))     dyf = DynamicFrame.fromDF(df, glueContext, "df")       return DynamicFrameCollection({"CustomTransform": dyf}, glueContext)
Without casting some MLlib column types to a String or integer first, you will get an error when trying to return the dataframe from the Custom Transform. VectorUDT is a ‘user-defined type’ of column that is often produced by MLlib transforms. Below is an example of casting the column to String before returning it.
# Cast VectorUDT column to StringrescaledData=rescaledData.withColumn("features",rescaledData["features"].cast("String"))# Reduce the number of columns being returnedrescaledData = rescaledData.selectExpr("labels", "features")# Convert from Spark dataframe to AWS Glue DynamicFramedyf_out = DynamicFrame.fromDF(rescaledData, glueContext, "rescaledData")   # Wrap the DynamicFrame in a DynamicFrameCollection and returnreturn DynamicFrameCollection({"CustomTransform": dyf_out}, glueContext)