fix: paginating describe log streams (#1046)

aws · Oct 9, 2019 · 5199b7e · 5199b7e
1 parent 5bba595
commit 5199b7e
Showing 1 changed file with 12 additions and 1 deletion.
diff --git a/src/sagemaker/session.py b/src/sagemaker/session.py
@@ -1985,9 +1985,20 @@ def _flush_log_streams(
                 logGroupName=log_group,
                 logStreamNamePrefix=job_name + "/",
                 orderBy="LogStreamName",
-                limit=instance_count,
+                limit=min(instance_count, 50),
             )
             stream_names = [s["logStreamName"] for s in streams["logStreams"]]
+
+            while "nextToken" in streams:
+                streams = client.describe_log_streams(
+                    logGroupName=log_group,
+                    logStreamNamePrefix=job_name + "/",
+                    orderBy="LogStreamName",
+                    limit=50,
+                )
+
+                stream_names.extend([s["logStreamName"] for s in streams["logStreams"]])
+
             positions.update(
                 [
                     (s, sagemaker.logs.Position(timestamp=0, skip=0))