fix: use al2 base image from ecr (#38)

* fix: use al2 from ecr

* extend timeout in job

Author: guoqiao1992

scripts/build.sh

@@ -22,10 +22,14 @@ source scripts/shared.sh
 
 parse_std_args "$@"
 
+aws ecr get-login-password --region us-west-2 | docker login --username AWS --password-stdin 137112412989.dkr.ecr.us-west-2.amazonaws.com
+
 echo "building image ${version} ... "
 docker build \
     -f ${build_context}/docker/Dockerfile.${processor} \
     -t ${repository}:${version} \
     --build-arg REGION=${REGION} \
     -t sagemaker-spark:latest \
     ${build_context}
+
+docker logout https://137112412989.dkr.ecr.us-west-2.amazonaws.com

spark/processing/3.0/py3/docker/Dockerfile.cpu

@@ -1,4 +1,4 @@
-FROM amazonlinux:2
+FROM 137112412989.dkr.ecr.us-west-2.amazonaws.com/amazonlinux:2
 ARG REGION
 ENV AWS_REGION ${REGION}
 RUN yum clean all

src/smspark/job.py

@@ -32,6 +32,9 @@
 class ProcessingJobManager(object):
     """Manages the lifecycle of a Spark job."""
 
+    _bootstrapping_timeout = 600.0  # all hosts should report as ready within this timeout.
+    _wait_for_primary_timeout = 600.0  # then, all workers ask the primary if it's up within this timeout.
+
     def __init__(
         self,
         resource_config: Dict[str, Any] = None,  # type: ignore
@@ -136,7 +139,11 @@ def all_hosts_have_bootstrapped() -> bool:
                 has_bootstrapped = [message.status == Status.WAITING for message in host_statuses.values()]
                 return all(has_bootstrapped)
 
-            self.waiter.wait_for(predicate_fn=all_hosts_have_bootstrapped, timeout=180.0, period=5.0)
+            self.waiter.wait_for(
+                predicate_fn=all_hosts_have_bootstrapped,
+                timeout=ProcessingJobManager._bootstrapping_timeout,
+                period=5.0,
+            )
 
             try:
                 subprocess.run(spark_submit_cmd, check=True, shell=True)
@@ -172,7 +179,7 @@ def primary_is_down() -> bool:
                 return not primary_is_up()
 
             self.logger.info("waiting for the primary to come up")
-            self.waiter.wait_for(primary_is_up, timeout=60.0, period=1.0)
+            self.waiter.wait_for(primary_is_up, timeout=ProcessingJobManager._wait_for_primary_timeout, period=1.0)
             self.logger.info("waiting for the primary to go down")
             self.waiter.wait_for(primary_is_down, timeout=float("inf"), period=5.0)
             self.logger.info("primary is down, worker now exiting")