Handle StopIteration in CloudWatch Logs retrieval during training

Author: yangaws

CHANGELOG.rst

@@ -2,6 +2,11 @@
 CHANGELOG
 =========
 
+1.17.1.dev
+==========
+
+* bug-fix: Handle StopIteration in CloudWatch Logs retrieval
+
 1.17.0
 ======
 

src/sagemaker/logs.py

@@ -94,7 +94,15 @@ def multi_stream_iter(client, log_group, streams, positions=None):
     """
     positions = positions or {s: Position(timestamp=0, skip=0) for s in streams}
     event_iters = [log_stream(client, log_group, s, positions[s].timestamp, positions[s].skip) for s in streams]
-    events = [next(s) if s else None for s in event_iters]
+    events = []
+    for s in event_iters:
+        if not s:
+            events.append(None)
+            continue
+        try:
+            events.append(next(s))
+        except StopIteration:
+            events.append(None)
 
     while some(events):
         i = argmin(events, lambda x: x['timestamp'] if x else 9999999999)