Minor fixes

aim/cli/server/commands.py

@@ -100,10 +100,10 @@ def server(host, port,
     click.echo('Press Ctrl+C to exit')
 
     try:
-        cmd = build_uvicorn_command('aim.ext.tracking.run:app',
+        cmd = build_uvicorn_command('aim.ext.transport.run:app',
                                     host=host, port=port,
                                     ssl_keyfile=ssl_keyfile, ssl_certfile=ssl_certfile, log_level=log_level)
-        exec_cmd(cmd)
+        exec_cmd(cmd, stream_output=True)
     except Exception:
         click.echo('Failed to run Aim Tracking Server. '
                    'Please see the logs above for details.')

aim/cli/up/commands.py

@@ -134,7 +134,7 @@ def up(dev, host, port, workers, uds,
 
     try:
         server_cmd = build_uvicorn_command(
-            'aim._ext.web.run:app',
+            'aim.web.run:app',
             host=host, port=port,
             workers=workers, uds=uds,
             ssl_keyfile=ssl_keyfile, ssl_certfile=ssl_certfile,

aim/ext/transport/client.py

@@ -134,6 +134,8 @@ def client_heartbeat(self):
 
         return response
 
+    @handle_exception(requests.ConnectionError,
+                      error_message='Failed to connect to Aim Server. Have you forgot to run `aim server` command?')
     def connect(self):
         endpoint = f'{self._http_protocol}{self._client_endpoint}/connect/{self.uri}/'
         response = requests.get(endpoint, headers=self.request_headers)

aim/ext/transport/server.py

@@ -43,10 +43,7 @@ def create_app():
         max_age=86400
     )
 
-    registry = ResourceTypeRegistry()
-    registry.register('TreeView', get_tree)
-    registry.register('FileManager', get_file_manager)
-    registry.register('Repo', get_repo)
+    registry = prepare_resource_registry()
 
     client_router = ClientRouter()
     tracking_router = TrackingRouter(registry)

aim/sdk/repo.py

@@ -79,6 +79,8 @@ def _close(self) -> None:
         """Close the `Repo` and unmount the remote repository."""
         if self._client:
             self._client._heartbeat_sender.stop()
+            self._client.get_queue().wait_for_finish()
+            self._client.get_queue().stop()
             self._client.disconnect()
         if self._mount_root:
             logger.debug(f'Unmounting remote repository at {self._mount_root}')
@@ -933,7 +935,7 @@ def is_remote_repo(self):
     @contextmanager
     def atomic_track(self, queue_id):
         if self.is_remote_repo:
-            self._client.start_instructions_batch()
+            self._client.start_instructions_batch(queue_id)
         yield
         if self.is_remote_repo:
             self._client.flush_instructions_batch(queue_id)

aim/sdk/run.py

@@ -91,12 +91,7 @@ def finalize_run(self):
 
     def empty_rpc_queue(self):
         if self.repo.is_remote_repo:
-            self.repo._client.get_queue(self.hash).wait_for_finish()
-
-    def finalize_rpc_queue(self):
-        if self.repo.is_remote_repo:
-            self.repo._client.get_queue(self.hash).stop()
-            self.repo._client.remove_queue(self.hash)
+            self.repo._client.get_queue().wait_for_finish()
 
     def _close(self) -> None:
         """
@@ -117,7 +112,6 @@ def _close(self) -> None:
             self._checkins.close()
         if self._lock:
             self._lock.release()
-        self.finalize_rpc_queue()
 
 
 # TODO: [AT] generate automatically based on ModelMappedRun

aim/web/utils.py

@@ -1,6 +1,5 @@
 from importlib import import_module
 import os
-import subprocess
 
 from aim.web.configs import AIM_UI_MOUNTED_REPO_PATH
 from aim.sdk.configs import get_aim_repo_name