Thread-safe core mlc operations: rm cache, mark-tmp, repos.json

mlc/action.py

@@ -415,15 +415,17 @@ def rm(self, i):
             item_meta = result.meta
 
             if os.path.exists(item_path):
-                if force_remove == True:
-                    shutil.rmtree(item_path)
-                else:
+                if force_remove != True:
                     user_choice = input(
                         f"Confirm to delete {target_name} item: {item_path}? (yes/no): ").strip().lower()
                     if user_choice not in ['yes', 'y']:
                         continue
-                    else:
-                        shutil.rmtree(item_path)
+
+                try:
+                    shutil.rmtree(item_path)
+                except FileNotFoundError:
+                    logger.warning(
+                        f"{item_path} was already removed by another process.")
 
                 logger.info(
                     f"{target_name} item: {item_path} has been successfully removed")

mlc/cache_action.py

@@ -4,6 +4,7 @@
 import time
 from . import utils
 from .logger import logger
+from filelock import FileLock, Timeout
 
 
 class CacheAction(Action):
@@ -135,22 +136,25 @@ def mark_tmp(self, i):
         updated_count = 0
 
         for item in res['list']:
-            tags = item.meta.get("tags", [])
-            if 'tmp' in tags:
-                continue
-
-            tags.append('tmp')
-            item.meta["tags"] = tags
             meta_yaml_path = os.path.join(item.path, "meta.yaml")
             meta_json_path = os.path.join(item.path, "meta.json")
+            lock_file = item.path + ".lock"
+
+            with FileLock(lock_file, timeout=60):
+                tags = item.meta.get("tags", [])
+                if 'tmp' in tags:
+                    continue
+
+                tags.append('tmp')
+                item.meta["tags"] = tags
 
-            if os.path.exists(meta_yaml_path):
-                save_result = utils.save_yaml(meta_yaml_path, meta=item.meta)
-            else:
-                save_result = utils.save_json(meta_json_path, meta=item.meta)
+                if os.path.exists(meta_yaml_path):
+                    save_result = utils.save_yaml(meta_yaml_path, meta=item.meta)
+                else:
+                    save_result = utils.save_json(meta_json_path, meta=item.meta)
 
-            if save_result['return'] > 0:
-                return save_result
+                if save_result['return'] > 0:
+                    return save_result
 
             self.get_index().update(item.meta, "cache", item.path, item.repo)
             updated_count += 1

mlc/repo_action.py

@@ -10,6 +10,7 @@
 from urllib.parse import urlparse
 from .repo import Repo
 from .index import Index
+from filelock import FileLock, Timeout
 
 
 class RepoAction(Action):
@@ -173,16 +174,17 @@ def register_repo(self, repo_path, repo_meta, ignore_on_conflict=False):
         # Get the path to the repos.json file in $HOME/MLC
         repos_file_path = os.path.join(self.repos_path, 'repos.json')
 
-        with open(repos_file_path, 'r') as f:
-            repos_list = json.load(f)
+        with FileLock(_repos_lock_file(repos_file_path), timeout=60):
+            with open(repos_file_path, 'r') as f:
+                repos_list = json.load(f)
 
-        if repo_path not in repos_list:
-            repos_list.append(repo_path)
-            logger.info(f"Added new repo path: {repo_path}")
+            if repo_path not in repos_list:
+                repos_list.append(repo_path)
+                logger.info(f"Added new repo path: {repo_path}")
 
-        with open(repos_file_path, 'w') as f:
-            json.dump(repos_list, f, indent=2)
-            logger.info(f"Updated repos.json at {repos_file_path}")
+            with open(repos_file_path, 'w') as f:
+                json.dump(repos_list, f, indent=2)
+                logger.info(f"Updated repos.json at {repos_file_path}")
 
         self.repos = self.load_repos_and_meta()
         repo_obj = next(
@@ -727,6 +729,11 @@ def rm(self, run_args):
         return rm_repo(repo_path, repos_file_path, force_remove)
 
 
+def _repos_lock_file(repos_file_path):
+    """Return the lock file path for a repos.json file."""
+    return repos_file_path + ".lock"
+
+
 def rm_repo(repo_path, repos_file_path, force_remove):
     logger.info(
         "rm command has been called for repo. This would delete the repo folder and unregister the repo from repos.json")
@@ -781,16 +788,17 @@ def rm_repo(repo_path, repos_file_path, force_remove):
 def unregister_repo(repo_path, repos_file_path):
     logger.info(f"Unregistering the repo in path {repo_path}")
 
-    with open(repos_file_path, 'r') as f:
-        repos_list = json.load(f)
+    with FileLock(_repos_lock_file(repos_file_path), timeout=60):
+        with open(repos_file_path, 'r') as f:
+            repos_list = json.load(f)
 
-    if repo_path in repos_list:
-        repos_list.remove(repo_path)
-        with open(repos_file_path, 'w') as f:
-            json.dump(repos_list, f, indent=2)
-        logger.info(f"Path: {repo_path} has been removed.")
-    else:
-        logger.info(
-            f"Path: {repo_path} not found in {repos_file_path}. Nothing to be unregistered!")
+        if repo_path in repos_list:
+            repos_list.remove(repo_path)
+            with open(repos_file_path, 'w') as f:
+                json.dump(repos_list, f, indent=2)
+            logger.info(f"Path: {repo_path} has been removed.")
+        else:
+            logger.info(
+                f"Path: {repo_path} not found in {repos_file_path}. Nothing to be unregistered!")
 
     return {'return': 0}

tests/test_thread_safety.py

@@ -0,0 +1,183 @@
+"""
+Thread safety tests for core mlc operations:
+- mlc rm cache
+- mlc mark-tmp cache
+- register_repo / unregister_repo (repos.json)
+"""
+
+import json
+import os
+import tempfile
+import threading
+import unittest
+
+from mlc.action import Action
+from mlc.cache_action import CacheAction
+from mlc.repo_action import unregister_repo
+
+
+class ConcurrentRmCacheTest(unittest.TestCase):
+    """Concurrent mlc rm cache calls must not raise or corrupt state."""
+
+    def setUp(self):
+        self.temp_dir = tempfile.TemporaryDirectory()
+        self.addCleanup(self.temp_dir.cleanup)
+        self.previous_mlc_repos = os.environ.get("MLC_REPOS")
+        self.addCleanup(self._restore_env)
+
+        repos_dir = os.path.join(self.temp_dir.name, "repos")
+        os.environ["MLC_REPOS"] = repos_dir
+
+        action = Action()
+        action.parent = None
+        for name, tags in [("cache-a", "get,dataset,a"), ("cache-b", "get,dataset,b")]:
+            res = action.add({"target_name": "cache", "item": name, "tags": tags})
+            self.assertEqual(res["return"], 0)
+
+    def _restore_env(self):
+        if self.previous_mlc_repos is None:
+            os.environ.pop("MLC_REPOS", None)
+        else:
+            os.environ["MLC_REPOS"] = self.previous_mlc_repos
+
+    def test_concurrent_rm_cache_does_not_raise(self):
+        """Two threads deleting the same cache item must not raise an exception."""
+        errors = []
+
+        def rm_cache(tags):
+            try:
+                action = Action()
+                action.parent = None
+                # Call rm directly with target_name set (same as CacheAction.rm does)
+                action.rm({"target_name": "cache", "tags": tags, "f": True})
+            except Exception as exc:
+                errors.append(exc)
+
+        t1 = threading.Thread(target=rm_cache, args=("get,dataset,a",))
+        t2 = threading.Thread(target=rm_cache, args=("get,dataset,a",))
+        t1.start()
+        t2.start()
+        t1.join()
+        t2.join()
+
+        self.assertEqual(errors, [], f"Unexpected exceptions: {errors}")
+
+
+class ConcurrentMarkTmpTest(unittest.TestCase):
+    """Concurrent mark-tmp calls on the same item must not duplicate the tag."""
+
+    def setUp(self):
+        self.temp_dir = tempfile.TemporaryDirectory()
+        self.addCleanup(self.temp_dir.cleanup)
+        self.previous_mlc_repos = os.environ.get("MLC_REPOS")
+        self.addCleanup(self._restore_env)
+
+        repos_dir = os.path.join(self.temp_dir.name, "repos")
+        os.environ["MLC_REPOS"] = repos_dir
+
+        action = Action()
+        action.parent = None
+        res = action.add({"target_name": "cache", "item": "shared-cache", "tags": "get,shared"})
+        self.assertEqual(res["return"], 0)
+        self.cache_path = res["path"]
+
+    def _restore_env(self):
+        if self.previous_mlc_repos is None:
+            os.environ.pop("MLC_REPOS", None)
+        else:
+            os.environ["MLC_REPOS"] = self.previous_mlc_repos
+
+    def test_concurrent_mark_tmp_no_duplicate_tag(self):
+        """Concurrent mark-tmp calls must not produce a duplicated 'tmp' tag."""
+        errors = []
+
+        def mark_tmp():
+            try:
+                action = Action()
+                action.parent = None
+                # CacheAction.__init__ does self.__dict__.update(vars(parent)),
+                # which would overwrite self.parent with parent.parent (None).
+                # Reassign explicitly so mark_tmp's self.search works via parent.
+                cache = CacheAction(action)
+                cache.parent = action
+                cache.mark_tmp({"tags": "get,shared"})
+            except Exception as exc:
+                errors.append(exc)
+
+        threads = [threading.Thread(target=mark_tmp) for _ in range(5)]
+        for t in threads:
+            t.start()
+        for t in threads:
+            t.join()
+
+        self.assertEqual(errors, [], f"Unexpected exceptions: {errors}")
+
+        meta_json = os.path.join(self.cache_path, "meta.json")
+        with open(meta_json) as f:
+            meta = json.load(f)
+        self.assertIn("tmp", meta["tags"])
+        self.assertEqual(meta["tags"].count("tmp"), 1,
+                         "Expected exactly one 'tmp' tag after concurrent mark-tmp calls")
+
+
+class ConcurrentReposJsonTest(unittest.TestCase):
+    """Concurrent unregister_repo calls must not corrupt repos.json."""
+
+    def setUp(self):
+        self.temp_dir = tempfile.TemporaryDirectory()
+        self.addCleanup(self.temp_dir.cleanup)
+        self.previous_mlc_repos = os.environ.get("MLC_REPOS")
+        self.addCleanup(self._restore_env)
+
+        self.repos_dir = os.path.join(self.temp_dir.name, "repos")
+        os.makedirs(self.repos_dir, exist_ok=True)
+        self.repos_file = os.path.join(self.repos_dir, "repos.json")
+
+        local_path = os.path.join(self.repos_dir, "local")
+        os.makedirs(local_path, exist_ok=True)
+        with open(os.path.join(local_path, "meta.yaml"), "w") as f:
+            f.write("uid: 0000000000000000\nalias: local\n")
+        with open(self.repos_file, "w") as f:
+            json.dump([local_path], f)
+
+        os.environ["MLC_REPOS"] = self.repos_dir
+
+    def _restore_env(self):
+        if self.previous_mlc_repos is None:
+            os.environ.pop("MLC_REPOS", None)
+        else:
+            os.environ["MLC_REPOS"] = self.previous_mlc_repos
+
+    def test_concurrent_unregister_no_corruption(self):
+        """Concurrent unregister_repo calls must leave repos.json as valid JSON
+        and must not lose any update due to a read-modify-write race."""
+        paths = [os.path.join(self.repos_dir, f"repo{i}") for i in range(10)]
+        with open(self.repos_file, "w") as f:
+            json.dump(paths, f)
+
+        errors = []
+
+        def unreg(p):
+            try:
+                unregister_repo(p, self.repos_file)
+            except Exception as exc:
+                errors.append(exc)
+
+        threads = [threading.Thread(target=unreg, args=(p,)) for p in paths]
+        for t in threads:
+            t.start()
+        for t in threads:
+            t.join()
+
+        self.assertEqual(errors, [], f"Unexpected exceptions: {errors}")
+
+        with open(self.repos_file) as f:
+            remaining = json.load(f)
+        self.assertIsInstance(remaining, list)
+        for p in paths:
+            self.assertNotIn(p, remaining,
+                             f"{p} should have been removed but was still in repos.json")
+
+
+if __name__ == "__main__":
+    unittest.main()