Pierre/UI improvements

aim/VERSION

@@ -1 +1 @@
-3.29.1
+3.30.0

aim/sdk/index_manager.py

@@ -134,19 +134,25 @@ def stop(self):
 
     def _monitor_existing_chunks(self):
         while not self._stop_event.is_set():
-            index_db = self.repo.request_tree('meta', read_only=True)
-            monitored_chunks = set(self._watches.keys())
-            for chunk_path in self.chunks_dir.iterdir():
-                if (
-                    chunk_path.is_dir()
-                    and chunk_path.name not in monitored_chunks
-                    and self._is_run_index_outdated(chunk_path.name, index_db)
-                ):
-                    logger.debug(f'Monitoring existing chunk: {chunk_path}')
-                    self.monitor_chunk_directory(chunk_path)
-                    logger.debug(f'Triggering indexing for run {chunk_path.name}')
-                    self.add_run_to_queue(chunk_path.name)
-            self.repo.container_pool.clear()
+            try:
+                index_db = self.repo.request_tree('meta', read_only=True)
+                monitored_chunks = set(self._watches.keys())
+                for chunk_path in self.chunks_dir.iterdir():
+                    try:
+                        if (
+                            chunk_path.is_dir()
+                            and chunk_path.name not in monitored_chunks
+                            and self._is_run_index_outdated(chunk_path.name, index_db)
+                        ):
+                            logger.debug(f'Monitoring existing chunk: {chunk_path}')
+                            self.monitor_chunk_directory(chunk_path)
+                            logger.debug(f'Triggering indexing for run {chunk_path.name}')
+                            self.add_run_to_queue(chunk_path.name)
+                    except Exception as e:
+                        logger.warning(f'Error checking chunk {chunk_path}: {e}')
+                self.repo.container_pool.clear()
+            except Exception as e:
+                logger.error(f'_monitor_existing_chunks iteration failed: {e}')
             time.sleep(5)
 
     def _stop_monitoring_chunk(self, run_hash):
@@ -155,9 +161,19 @@ def _stop_monitoring_chunk(self, run_hash):
             self.chunk_change_observer.unschedule(watch)
             logger.debug(f'Stopped monitoring chunk: {run_hash}')
 
+    # Maximum number of chunk directories to watch simultaneously.
+    # PollingObserver opens file descriptors for each watch; capping this
+    # prevents "Too many open files" when there are hundreds of runs.
+    MAX_WATCHED_CHUNKS = 50
+
     def monitor_chunk_directory(self, chunk_path):
         """Ensure chunk directory is monitored using a single handler."""
         if chunk_path.name not in self._watches:
+            if len(self._watches) >= self.MAX_WATCHED_CHUNKS:
+                # Drop the oldest watch to stay under the fd limit.
+                oldest = next(iter(self._watches))
+                self._stop_monitoring_chunk(oldest)
+                logger.debug(f'Watch limit reached, dropped oldest watch: {oldest}')
             watch = self.chunk_change_observer.schedule(self.chunk_change_handler, chunk_path, recursive=True)
             self._watches[chunk_path.name] = watch
             logger.debug(f'Started monitoring chunk directory: {chunk_path}')
@@ -176,12 +192,19 @@ def _process_indexing_queue(self):
         while not self._stop_event.is_set():
             _, run_hash = self.indexing_queue.get()
             logger.debug(f'Indexing run {run_hash}...')
-            self.index(run_hash)
-            self.indexing_queue.task_done()
+            try:
+                self.index(run_hash)
+            except Exception as e:
+                # An unhandled exception here would silently kill this daemon
+                # thread, leaving the indexing queue permanently stalled.
+                logger.error(f'Unexpected error indexing run {run_hash}: {e}')
+            finally:
+                self.indexing_queue.task_done()
 
     def index(self, run_hash):
-        index = self.repo._get_index_tree('meta', 0).view(())
+        import gc
         try:
+            index = self.repo._get_index_tree('meta', 0).view(())
             run_checksum = self._get_run_checksum(run_hash)
             meta_tree = self.repo.request_tree('meta', run_hash, read_only=True, skip_read_optimization=True).subtree(
                 'meta'
@@ -197,6 +220,19 @@ def index(self, run_hash):
         except (aimrocks.errors.RocksIOError, aimrocks.errors.Corruption):
             logger.warning(f'Indexing thread detected corrupted run: {run_hash}. Skipping.')
             self._corrupted_runs.add(run_hash)
+            self._stop_monitoring_chunk(run_hash)
+        except Exception as e:
+            # Catch-all: log and skip rather than propagating to
+            # _process_indexing_queue where it would kill the thread.
+            logger.warning(f'Indexing run {run_hash} failed unexpectedly: {e}. Skipping.')
+            self._corrupted_runs.add(run_hash)
+            self._stop_monitoring_chunk(run_hash)
+        finally:
+            # Release TreeView references so RocksDB containers can be GC'd
+            # promptly. Without this, WeakValueDictionary entries stay alive
+            # until the next GC cycle, accumulating open file descriptors.
+            self.repo.container_pool.clear()
+            gc.collect()
         return True
 
     def _is_run_index_outdated(self, run_hash, index_db):

aim/sdk/lock_manager.py

@@ -20,6 +20,10 @@
 
 logger = logging.getLogger(__name__)
 
+# Locks written by a different machine cannot have their PID checked remotely.
+# Treat them as stale after this many hours of inactivity.
+STALE_LOCK_THRESHOLD_HOURS = 4
+
 
 class LockingVersion(Enum):
     LEGACY = 0
@@ -164,8 +168,21 @@ def release_locks(self, run_hash: str, force: bool) -> bool:
         return success
 
     def is_stalled_lock(self, lock_file_path: Path) -> bool:
-        with open(lock_file_path, mode='r') as lock_metadata_fh:
-            machine_id, pid, *_ = lock_metadata_fh.read().split('-')
-            if int(machine_id) == self.machine_id and not psutil.pid_exists(int(pid)):
-                return True
-        return False
+        try:
+            with open(lock_file_path, mode='r') as lock_metadata_fh:
+                parts = lock_metadata_fh.read().strip().split('-')
+                machine_id, pid = int(parts[0]), int(parts[1])
+            if machine_id == self.machine_id:
+                # Same machine: check if the owning process is still alive.
+                return not psutil.pid_exists(pid)
+            # Cross-machine lock: the remote PID cannot be checked from here.
+            # Fall back to file age — if the lock file has not been refreshed
+            # in STALE_LOCK_THRESHOLD_HOURS it is almost certainly abandoned
+            # (e.g. a training node that was killed with SIGKILL or SIGTERM
+            # without releasing the lock).
+            age = datetime.datetime.now() - datetime.datetime.fromtimestamp(
+                lock_file_path.stat().st_mtime
+            )
+            return age > datetime.timedelta(hours=STALE_LOCK_THRESHOLD_HOURS)
+        except Exception:
+            return False

aim/sdk/repo.py

@@ -397,10 +397,11 @@ def list_corrupted_runs(self) -> List[str]:
         from aim.storage.encoding import decode_path
 
         def get_run_hash_from_prefix(prefix: bytes):
-            return decode_path(prefix)[-1]
+            parts = decode_path(prefix)
+            return parts[-1] if parts else None
 
         container = RocksUnionContainer(os.path.join(self.path, 'meta'), read_only=True)
-        return list(map(get_run_hash_from_prefix, container.corrupted_dbs))
+        return [h for h in map(get_run_hash_from_prefix, container.corrupted_dbs) if h is not None]
 
     def _active_run_hashes(self) -> Set[str]:
         if self.is_remote_repo:
@@ -806,7 +807,9 @@ def _delete_experiment(self, exp_id):
 
     def _delete_local_run_data(self, run_hash: str):
         # remove data from index container
-        index_tree = self._get_index_container('meta', timeout=0).tree()
+        # timeout=30: the index daemon holds the LOCK continuously; give it
+        # enough time to finish its current write before we acquire it.
+        index_tree = self._get_index_container('meta', timeout=30).tree()
         del index_tree.subtree(('meta', 'chunks'))[run_hash]
 
         # delete rocksdb containers data
@@ -963,7 +966,7 @@ def optimize_container(path, extra_options):
             rc.optimize_for_read()
 
         if self.is_remote_repo:
-            self._remote_repo_proxy._close_run(run_hash)
+            return self._remote_repo_proxy._close_run(run_hash)
 
         lock_manager = LockManager(self.path)
 
@@ -978,6 +981,13 @@ def optimize_container(path, extra_options):
             if not meta_run_tree.get('end_time'):
                 meta_run_tree['end_time'] = datetime.datetime.now(pytz.utc).timestamp()
 
+            # Remove the progress file so list_active_runs() no longer returns
+            # this run. Without this, a crashed run stays in meta/progress/ forever
+            # and the active-runs streamer keeps polling it on every UI refresh.
+            progress_path = os.path.join(self.path, 'meta', 'progress', run_hash)
+            if os.path.exists(progress_path):
+                os.remove(progress_path)
+
             # Run rocksdb optimizations if container locks are removed
             meta_db_path = os.path.join(self.path, 'meta', 'chunks', run_hash)
             seqs_db_path = os.path.join(self.path, 'seqs', 'chunks', run_hash)

aim/sdk/reporter/__init__.py

@@ -804,3 +804,8 @@ def _run(self):
     def stop(self):
         self.stop_signal.set()
         self.thread.join()
+        if self.touch_path is not None and self.touch_path.exists():
+            try:
+                self.touch_path.unlink()
+            except OSError:
+                pass

aim/web/api/runs/object_api_utils.py

@@ -25,10 +25,15 @@
 
 
 def get_blobs_batch(uri_batch: List[str], repo: 'Repo') -> Iterator[bytes]:
+    import logging as _logging
+    _logger = _logging.getLogger(__name__)
     uri_service = URIService(repo=repo)
     batch_iterator = uri_service.request_batch(uri_batch=uri_batch)
-    for it in batch_iterator:
-        yield collect_streamable_data(encode_tree(it))
+    try:
+        for it in batch_iterator:
+            yield collect_streamable_data(encode_tree(it))
+    except Exception as e:
+        _logger.warning(f'get_blobs_batch: skipping blob due to error: {e}')
 
 
 class CustomObjectApi:
@@ -143,10 +148,17 @@ def _pack_run_data(run_: 'Run', traces_: list):
                     progress_reports_sent += 1
                     last_reported_progress_time = time.time()
                 if run_info.get('traces') and run_info.get('run'):
-                    traces_list = []
-                    for trace in run_info['traces']:
-                        traces_list.append(self._get_trace_info(trace, True, True))
-                    yield _pack_run_data(run_info['run'], traces_list)
+                    try:
+                        traces_list = []
+                        for trace in run_info['traces']:
+                            traces_list.append(self._get_trace_info(trace, True, True))
+                        yield _pack_run_data(run_info['run'], traces_list)
+                    except Exception as _e:
+                        import logging as _logging
+                        _logging.getLogger(__name__).warning(
+                            f'search_result_streamer: skipping run {run_info["run"].hash} due to error: {_e}'
+                        )
+                        continue
                     if report_progress:
                         yield collect_streamable_data(
                             encode_tree({f'progress_{progress_reports_sent}': run_info['progress']})
@@ -162,18 +174,23 @@ def _pack_run_data(run_: 'Run', traces_: list):
             pass
 
     async def requested_traces_streamer(self) -> List[dict]:
+        import logging as _logging
+        _logger = _logging.getLogger(__name__)
         try:
             for key in list(self.trace_cache.keys()):
                 run_info = self.trace_cache[key]
                 await asyncio.sleep(ASYNC_SLEEP_INTERVAL)
                 for trace in run_info['traces']:
-                    trace_dict = self._get_trace_info(trace, False, False)
-                    trace_dict['record_range_used'] = self.record_range
-                    trace_dict['record_range_total'] = self.total_record_range
-                    if self.use_list:
-                        trace_dict['index_range'] = self.index_range
-                        trace_dict['index_range_total'] = self.total_index_range
-                    yield collect_streamable_data(encode_tree(trace_dict))
+                    try:
+                        trace_dict = self._get_trace_info(trace, False, False)
+                        trace_dict['record_range_used'] = self.record_range
+                        trace_dict['record_range_total'] = self.total_record_range
+                        if self.use_list:
+                            trace_dict['index_range'] = self.index_range
+                            trace_dict['index_range_total'] = self.total_index_range
+                        yield collect_streamable_data(encode_tree(trace_dict))
+                    except Exception as _e:
+                        _logger.warning(f'requested_traces_streamer: skipping trace due to error: {_e}')
                 del self.trace_cache[key]
             self.run = None
         except asyncio.CancelledError: