Revert "Fix log messages (#549)"

This reverts commit 2b229cd.

CHANGELOG.md

@@ -14,8 +14,6 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Removed
 
-- Improve display of stream log messages. ([#549](https://github.com/jina-ai/finetuner/pull/549))
-
 ### Changed
 
 - Bump hubble client version. ([#546](https://github.com/jina-ai/finetuner/pull/546))

finetuner/__init__.py

@@ -4,6 +4,8 @@
 from typing import Any, Dict, List, Optional, Union
 
 from docarray import DocumentArray
+from rich.console import Console
+from rich.table import Table
 from stubs import model as model_stub
 
 from finetuner.constants import (
@@ -22,10 +24,8 @@
     os.environ[HUBBLE_REGISTRY] = DEFAULT_HUBBLE_REGISTRY
 
 from finetuner import callback, model
-from finetuner.console import print_model_table
 from finetuner.experiment import Experiment
 from finetuner.finetuner import Finetuner
-from finetuner.model import list_model_classes
 
 ft = Finetuner()
 
@@ -45,6 +45,15 @@ def list_callbacks() -> Dict[str, callback.CallbackStubType]:
     }
 
 
+def _list_models() -> Dict[str, model.ModelStubType]:
+    rv = {}
+    members = inspect.getmembers(model_stub, inspect.isclass)
+    for name, stub in members:
+        if name != 'MLPStub' and not name.startswith('_') and type(stub) != type:
+            rv[name] = stub
+    return rv
+
+
 def _build_name_stub_map() -> Dict[str, model.ModelStubType]:
     rv = {}
     members = inspect.getmembers(model_stub, inspect.isclass)
@@ -56,7 +65,7 @@ def _build_name_stub_map() -> Dict[str, model.ModelStubType]:
 
 def list_models() -> List[str]:
     """List available models for training."""
-    return [name for name in list_model_classes()]
+    return [name for name in _list_models()]
 
 
 def list_model_options() -> Dict[str, List[Dict[str, Any]]]:
@@ -76,13 +85,26 @@ def list_model_options() -> Dict[str, List[Dict[str, Any]]]:
             ).parameters.values()
             if parameter.name != 'self'
         ]
-        for name, _model_class in list_model_classes().items()
+        for name, _model_class in _list_models().items()
     }
 
 
 def describe_models() -> None:
     """Describe available models in a table."""
-    print_model_table(model)
+    table = Table(title='Finetuner backbones')
+    header = model.get_header()
+    model_descriptors = set()
+
+    for column in header:
+        table.add_column(column, justify='right', style='cyan', no_wrap=False)
+
+    for _, _model_class in _list_models().items():
+        if _model_class.descriptor not in model_descriptors:
+            table.add_row(*model.get_row(_model_class))
+            model_descriptors.add(_model_class.descriptor)
+
+    console = Console()
+    console.print(table)
 
 
 @login_required

finetuner/client/client.py

@@ -191,15 +191,7 @@ def stream_run_logs(self, experiment_name: str, run_name: str) -> Iterator[str]:
         response = self._handle_request(url=url, method=GET, stream=True)
         for entry in response.iter_lines():
             if entry:
-                decoded_message: str = entry.decode('utf-8', errors='ignore')
-                sep_pos = decoded_message.find(': ')
-                if sep_pos != -1:
-                    msg_type, msg = (
-                        decoded_message[:sep_pos],
-                        decoded_message[sep_pos + 2 :],
-                    )
-                    if msg_type in ('data', 'event'):
-                        yield msg
+                yield entry.decode('utf-8', errors='ignore')
 
     def create_run(
         self,

finetuner/model.py

@@ -1,6 +1,3 @@
-import inspect
-
-from stubs import model
 from stubs.model import *  # noqa F401
 
 
@@ -18,12 +15,3 @@ def get_row(model_stub) -> Tuple[str, ...]:
         model_stub.architecture,
         model_stub.description,
     )
-
-
-def list_model_classes() -> Dict[str, ModelStubType]:
-    rv = {}
-    members = inspect.getmembers(model, inspect.isclass)
-    for name, stub in members:
-        if name != 'MLPStub' and not name.startswith('_') and type(stub) != type:
-            rv[name] = stub
-    return rv

finetuner/run.py

@@ -2,7 +2,6 @@
 from typing import Iterator
 
 from finetuner.client import FinetunerV1Client
-from finetuner.console import console
 from finetuner.constants import (
     ARTIFACT_ID,
     ARTIFACTS_DIR,
@@ -82,17 +81,17 @@ def stream_logs(self, interval: int = 5) -> Iterator[str]:
         :param interval: The time interval to sync the status of finetuner `Run`.
         :yield: An iterators keep stream the logs from server.
         """
-        status = self.status()[STATUS]
-        msg_template = (
-            'Preparing to run, logs will be ready to pull when '
-            '`status` is `STARTED`. Current status is `%s`'
-        )
-        with console.status(msg_template % status, spinner="dots") as rich_status:
-            while status == CREATED:
-                time.sleep(interval)
-                status = self.status()[STATUS]
-                rich_status.update(msg_template % status)
-
+        while True:
+            time.sleep(interval)
+            status = self.status()[STATUS]
+            if status == CREATED:
+                msg = (
+                    f'Preparing to run, logs will be ready to pull when '
+                    f'`status` is `STARTED`. Current status is `{status}`'
+                )
+                print(msg)
+            else:
+                break
         return self._client.stream_run_logs(
             experiment_name=self._experiment_name, run_name=self._name
         )