[UsageStats] Record usage of ray client and ray job submission (#29089)

- Record usage of ray client and ray job submission - Better usage detection for tune: instead of import, record in tune.run. (Our dashboard will import tune which is a false positive) Signed-off-by: Jiajun Yao <[email protected]>
ray-project · Oct 6, 2022 · 889f769 · 889f769
1 parent 260dce9
commit 889f769
Show file tree

Hide file tree

Showing 6 changed files with 47 additions and 9 deletions.
diff --git a/dashboard/modules/job/job_head.py b/dashboard/modules/job/job_head.py
@@ -201,6 +201,8 @@ async def submit_job(self, req: Request) -> Response:
         request_submission_id = submit_request.submission_id or submit_request.job_id
 
         try:
+            ray._private.usage.usage_lib.record_library_usage("job_submission")
+
             submission_id = await self._job_manager.submit_job(
                 entrypoint=submit_request.entrypoint,
                 submission_id=request_submission_id,

diff --git a/python/ray/_private/usage/usage_lib.py b/python/ray/_private/usage/usage_lib.py
@@ -259,7 +259,7 @@ class TagKey(Enum):
     # The total number of running serve deployments as a string.
     SERVE_NUM_DEPLOYMENTS = auto()
 
-    # The GCS storage type, which could be memory or redis
+    # The GCS storage type, which could be memory or redis.
     GCS_STORAGE = auto()
 
 

diff --git a/python/ray/_private/worker.py b/python/ray/_private/worker.py
@@ -1260,6 +1260,8 @@ def init(
                 usage_lib.put_pre_init_usage_stats()
         else:
             usage_lib.put_pre_init_usage_stats()
+
+        usage_lib.record_library_usage("client")
         return ctx
 
     if kwargs:

diff --git a/python/ray/tests/test_usage_stats.py b/python/ray/tests/test_usage_stats.py
@@ -591,6 +591,16 @@ def get_actor_metadata(self):
         "ray://127.0.0.1:10001" if ray_client else address
     )
     run_string_as_driver(driver)
+
+    job_submission_client = ray.job_submission.JobSubmissionClient(
+        "http://127.0.0.1:8265"
+    )
+    job_id = job_submission_client.submit_job(entrypoint="ls")
+    wait_for_condition(
+        lambda: job_submission_client.get_job_status(job_id)
+        == ray.job_submission.JobStatus.SUCCEEDED
+    )
+
     library_usages = ray_usage_lib.get_library_usages_to_report(
         ray.experimental.internal_kv.internal_kv_get_gcs_client()
     )
@@ -609,7 +619,10 @@ def get_actor_metadata(self):
         "util.multiprocessing.Pool",
         "util.Queue",
         "util.joblib",
+        "job_submission",
     }
+    if ray_client:
+        expected.add("client")
     assert set(library_usages) == expected
     if not ray_client:
         assert set(lib_usages_from_home_folder) == expected
@@ -927,7 +940,6 @@ def test_usage_report_e2e(
         cluster.add_node(num_cpus=3)
         if os.environ.get("RAY_MINIMAL") != "1":
             from ray import train  # noqa: F401
-            from ray import tune  # noqa: F401
             from ray.rllib.algorithms.ppo import PPO  # noqa: F401
 
         ray_usage_lib.record_extra_usage_tag(ray_usage_lib.TagKey._TEST1, "extra_v2")
@@ -936,6 +948,15 @@ def test_usage_report_e2e(
 
         ray_usage_lib.record_extra_usage_tag(ray_usage_lib.TagKey._TEST2, "extra_v3")
 
+        if os.environ.get("RAY_MINIMAL") != "1":
+            from ray import tune  # noqa: F401
+
+            def objective(*args):
+                pass
+
+            tuner = tune.Tuner(objective)
+            tuner.fit()
+
         @ray.remote(num_cpus=0)
         class StatusReporter:
             def __init__(self):
@@ -1168,10 +1189,16 @@ def test_lib_used_from_driver(monkeypatch, ray_start_cluster, reset_usage_stats)
 import os
 if os.environ.get("RAY_MINIMAL") != "1":
     from ray import train  # noqa: F401
-    from ray import tune  # noqa: F401
     from ray.rllib.algorithms.ppo import PPO  # noqa: F401
 
 ray.init(address="{addr}")
+
+if os.environ.get("RAY_MINIMAL") != "1":
+    from ray import tune  # noqa: F401
+    def objective(*args):
+        pass
+
+    tune.run(objective)
 """
         # Run a script in a separate process. It is a workaround to
         # reimport libraries. Without this, `import train`` will become
@@ -1221,11 +1248,15 @@ def test_lib_used_from_workers(monkeypatch, ray_start_cluster, reset_usage_stats
         class ActorWithLibImport:
             def __init__(self):
                 from ray import train  # noqa: F401
-                from ray import tune  # noqa: F401
                 from ray.rllib.algorithms.ppo import PPO  # noqa: F401
 
             def ready(self):
-                pass
+                from ray import tune  # noqa: F401
+
+                def objective(*args):
+                    pass
+
+                tune.run(objective)
 
         # Use a runtime env to run tests in minimal installation.
         a = ActorWithLibImport.options(
@@ -1270,6 +1301,11 @@ def test_lib_usage_record_from_init_session(
 
 # Start a instance that disables usage stats.
 ray.init()
+
+def objective(*args):
+    pass
+
+tune.run(objective)
 """
 
     run_string_as_driver(script)

diff --git a/python/ray/tune/__init__.py b/python/ray/tune/__init__.py
@@ -44,10 +44,6 @@
 from ray.tune.tuner import Tuner
 from ray.tune.tune_config import TuneConfig
 
-from ray._private.usage import usage_lib
-
-usage_lib.record_library_usage("tune")
-
 __all__ = [
     "Trainable",
     "Callback",

diff --git a/python/ray/tune/tune.py b/python/ray/tune/tune.py
@@ -429,6 +429,8 @@ def _handle_string_queue():
 
     del remote_run_kwargs
 
+    ray._private.usage.usage_lib.record_library_usage("tune")
+
     all_start = time.time()
 
     if mode and mode not in ["min", "max"]: