diff --git a/dvc/config.py b/dvc/config.py index e7b89d787..d3f31e9d8 100644 --- a/dvc/config.py +++ b/dvc/config.py @@ -103,6 +103,7 @@ class RelPath(str): REMOTE_COMMON = { "url": str, "checksum_jobs": All(Coerce(int), Range(1)), + "jobs": All(Coerce(int), Range(1)), Optional("no_traverse"): Bool, # obsoleted "verify": Bool, } diff --git a/dvc/tree/base.py b/dvc/tree/base.py index 4f594ec32..bbcd81dae 100644 --- a/dvc/tree/base.py +++ b/dvc/tree/base.py @@ -77,6 +77,14 @@ class BaseTree: self.verify = config.get("verify", self.DEFAULT_VERIFY) self.path_info = None + @cached_property + def jobs(self): + return ( + self.config.get("jobs") + or (self.repo and self.repo.config["core"].get("jobs")) + or self.JOBS + ) + @cached_property def hash_jobs(self): return ( @@ -404,7 +412,7 @@ class BaseTree: dir_mode=dir_mode, ) ) - with ThreadPoolExecutor(max_workers=self.JOBS) as executor: + with ThreadPoolExecutor(max_workers=self.jobs) as executor: futures = [ executor.submit(download_files, from_info, to_info) for from_info, to_info in zip(from_infos, to_infos) diff --git a/tests/unit/remote/test_remote.py b/tests/unit/remote/test_remote.py index 3769fb2c3..ff296dd60 100644 --- a/tests/unit/remote/test_remote.py +++ b/tests/unit/remote/test_remote.py @@ -16,6 +16,17 @@ def test_remote_with_hash_jobs(dvc): assert tree.hash_jobs == 100 +def test_remote_with_jobs(dvc): + dvc.config["remote"]["with_jobs"] = { + "url": "s3://bucket/name", + "jobs": 100, + } + dvc.config["core"]["jobs"] = 200 + + tree = get_cloud_tree(dvc, name="with_jobs") + assert tree.jobs == 100 + + def test_remote_without_hash_jobs(dvc): dvc.config["remote"]["without_hash_jobs"] = {"url": "s3://bucket/name"} dvc.config["core"]["checksum_jobs"] = 200