From 2f620e9226f6a337cd04da9748b4f40cd19efd1f Mon Sep 17 00:00:00 2001
From: Ilya Pekelny <ilya@launchpadrecruits.com>
Date: Fri, 13 Nov 2020 21:30:02 +0200
Subject: [PATCH] 4504: configuration-based job limit. (#4872)

* Job count expected in a core config

Besides defaulted 4 * core_count number of workers sourced in the
"tree", the configuration file based value expected and prioritised.

* Multiply the configured value by cpu cores count.

* Don't multiply by the core count twice.

* A couple more usages and a question note.

* Don't multiply a number of jobs to a cpu count

Number of jobs set in a config file directly used to spawn workers.

* Remove the obsolete note.

* Disallow 0 jobs set in a config

* `jobs` config-based value is optional

* Apply suggestions from code review

Co-authored-by: Ruslan Kuprieiev <kupruser@gmail.com>

* CPU count based default

* Update dvc/tree/base.py

* Happy path unit test

* Update tests/unit/remote/test_remote.py

* Fix the accidentally removed blank line

* Update dvc/tree/base.py

Co-authored-by: Ruslan Kuprieiev <kupruser@gmail.com>
---
 dvc/config.py                    |  1 +
 dvc/tree/base.py                 | 10 +++++++++-
 tests/unit/remote/test_remote.py | 11 +++++++++++
 3 files changed, 21 insertions(+), 1 deletion(-)

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