Merge pull request #54 from InfuseAI/feature/sc-21679/primehub-sdk-make-the-job-stream-follow-print

Fix the log follow issue. Should show carriage return correctly

Author: qrtt1

primehub/deployments.py

@@ -401,7 +401,7 @@ def delete(self, id, **kwargs):
     # TODO: handle invalid pod
     @cmd(name='logs', description='Get deployment logs by id',
          optionals=[('pod', str), ('follow', toggle_flag), ('tail', int)])
-    def logs(self, id, **kwargs) -> Iterator[str]:
+    def logs(self, id, **kwargs) -> Iterator[bytes]:
         """
         Get logs of a deployment
 

primehub/jobs.py

@@ -362,7 +362,7 @@ def wait(self, id, **kwargs):
         return self.get(id)
 
     @cmd(name='logs', description='Get job logs by id', optionals=[('follow', toggle_flag), ('tail', int)])
-    def logs(self, id, **kwargs) -> Iterator[str]:
+    def logs(self, id, **kwargs) -> Iterator[bytes]:
         """
         Get logs of a job
 

primehub/utils/display.py

@@ -84,7 +84,7 @@ def wrapper_generator(gen):
     is_type_sent = False
     for x in gen:
         if not is_type_sent:
-            yield isinstance(x, str)
+            yield isinstance(x, (str, bytes))
             is_type_sent = True
         yield x
 
@@ -125,6 +125,8 @@ def display_single(self, action: dict, value: Any, file: TextIO):
         logger.debug('display-single')
         if isinstance(value, str):
             print(value, file=file)
+        elif isinstance(value, bytes):
+            file.write(value.decode())
         else:
             json.dump(value, file)
 
@@ -182,5 +184,7 @@ def display_single(self, action: dict, value: Any, file: TextIO):
         logger.debug('display-single')
         if isinstance(value, str):
             print(value, file=file)
+        elif isinstance(value, bytes):
+            file.write(value.decode())
         else:
             display_tree_like_format(value, file)

primehub/utils/http_client.py

@@ -37,7 +37,7 @@ def request(self, variables: dict, query: str, error_handler: Callable = None):
         except BaseException as e:
             raise RequestException(e)
 
-    def request_logs(self, endpoint, follow, tail) -> Iterator[str]:
+    def request_logs(self, endpoint, follow, tail) -> Iterator[bytes]:
         params = {'follow': 'false'}
         if follow:
             params['follow'] = 'true'
@@ -46,11 +46,8 @@ def request_logs(self, endpoint, follow, tail) -> Iterator[str]:
         headers = {'authorization': 'Bearer {}'.format(self.primehub_config.api_token)}
 
         with requests.get(endpoint, headers=headers, params=params, stream=follow) as response:
-            if follow:
-                for line in response.iter_lines():
-                    yield line.decode()
-            else:
-                yield response.text
+            for chunk in response.iter_content(chunk_size=8192):
+                yield chunk
 
     def request_file(self, endpoint, dest):
         headers = {'authorization': 'Bearer {}'.format(self.primehub_config.api_token)}

tests/test_http_logs.py

@@ -27,4 +27,4 @@ def test_request_logs(self):
         count = 0
         for x in g:
             count = count + 1
-        self.assertEqual(1, count, 'get one result when no following')
+        self.assertTrue(count > 1, 'Generator gives lots of lines')