Repository: iterative/datachain

Issue:
save metrics in real-time
### Description

`metrics.set()` API saves metrics in memory, and later sends it to the main process to save it to the database.

https://github.com/iterative/datachain/blob/af5e9cbcefefe5cde9e0528cf38b2f8ac810d69c/src/datachain/cli.py#L868

This should happen in real-time, inside the user's script.

And this API should also work for when scripts are run outside of `datachain query` command. 


Part of #360.





### Version Info

_No response_

Discussion:
cloudflare-pages[bot]: ## Deploying datachain-documentation with &nbsp;<a href="https://pages.dev"><img alt="Cloudflare Pages" src="https://user-images.githubusercontent.com/23264/106598434-9e719e00-654f-11eb-9e59-6167043cfa01.png" width="16"></a> &nbsp;Cloudflare Pages

<table><tr><td><strong>Latest commit:</strong> </td><td>
<code>992a807</code>
</td></tr>
<tr><td><strong>Status:</strong></td><td>&nbsp;✅&nbsp; Deploy successful!</td></tr>
<tr><td><strong>Preview URL:</strong></td><td>
<a href='https://9a163cbe.datachain-documentation.pages.dev'>https://9a163cbe.datachain-documentation.pages.dev</a>
</td></tr>
<tr><td><strong>Branch Preview URL:</strong></td><td>
<a href='https://save-metrics-realtime.datachain-documentation.pages.dev'>https://save-metrics-realtime.datachain-documentation.pages.dev</a>
</td></tr>
</table>

[View logs](https://dash.cloudflare.com/?to=/:account/pages/view/datachain-documentation/9a163cbe-f9c8-4d68-9561-8d6c23b50aaf)

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/iterative/datachain/pull/387?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iterative) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 86.73%. Comparing base [(`16ed477`)](https://app.codecov.io/gh/iterative/datachain/commit/16ed477e5ce45906393f3c25e8100059dab45e28?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iterative) to head [(`992a807`)](https://app.codecov.io/gh/iterative/datachain/commit/992a80713173cba7fcfb5710b9d62c4a239d4cb2?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iterative).

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main     #387      +/-   ##
==========================================
- Coverage   86.78%   86.73%   -0.06%     
==========================================
  Files          92       92              
  Lines       10121    10122       +1     
  Branches     2054     1688     -366     
==========================================
- Hits         8784     8779       -5     
- Misses        988      994       +6     
  Partials      349      349              
```

| [Flag](https://app.codecov.io/gh/iterative/datachain/pull/387/flags?src=pr&el=flags&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iterative) | Coverage Δ | |
|---|---|---|
| [datachain](https://app.codecov.io/gh/iterative/datachain/pull/387/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iterative) | `86.68% <100.00%> (-0.04%)` | :arrow_down: |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iterative#carryforward-flags-in-the-pull-request-comment) to find out more.

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/iterative/datachain/pull/387?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iterative).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iterative).


Patch:
diff --git a/src/datachain/catalog/catalog.py b/src/datachain/catalog/catalog.py
index 073c7f6..1f540ca 100644
--- a/src/datachain/catalog/catalog.py
+++ b/src/datachain/catalog/catalog.py
@@ -157,7 +157,6 @@ class QueryResult(NamedTuple):
     version: Optional[int]
     output: str
     preview: Optional[list[dict]]
-    metrics: dict[str, Any]
 
 
 class DatasetRowsFetcher(NodesThreadPool):
@@ -1973,7 +1972,6 @@ class Catalog:
             version=version,
             output=output,
             preview=exec_result.preview,
-            metrics=exec_result.metrics,
         )
 
     def run_query(
diff --git a/src/datachain/cli.py b/src/datachain/cli.py
index 691d8f1..7011613 100644
--- a/src/datachain/cli.py
+++ b/src/datachain/cli.py
@@ -864,8 +864,7 @@ def query(
             error_stack=error_stack,
         )
         raise
-
-    catalog.metastore.set_job_status(job_id, JobStatus.COMPLETE, metrics=result.metrics)
+    catalog.metastore.set_job_status(job_id, JobStatus.COMPLETE)
 
     show_records(result.preview, collapse_columns=not no_collapse)
 
diff --git a/src/datachain/query/dataset.py b/src/datachain/query/dataset.py
index 8c80762..502a178 100644
--- a/src/datachain/query/dataset.py
+++ b/src/datachain/query/dataset.py
@@ -60,7 +60,6 @@ from datachain.utils import (
     get_datachain_executable,
 )
 
-from .metrics import metrics
 from .schema import C, UDFParamSpec, normalize_param
 from .session import Session
 from .udf import UDFBase, UDFClassWrapper, UDFFactory, UDFType
@@ -1729,7 +1728,6 @@ def _get_output_fd_for_write() -> Union[str, int]:
 class ExecutionResult:
     preview: list[dict] = attrs.field(factory=list)
     dataset: Optional[tuple[str, int]] = None
-    metrics: dict[str, Any] = attrs.field(factory=dict)
 
 
 def _send_result(dataset_query: DatasetQuery) -> None:
@@ -1766,7 +1764,7 @@ def _send_result(dataset_query: DatasetQuery) -> None:
         dataset = dataset_query.name, dataset_query.version
 
     preview = preview_query.to_db_records()
-    result = ExecutionResult(preview, dataset, metrics)
+    result = ExecutionResult(preview, dataset)
     data = attrs.asdict(result)
 
     with open(_get_output_fd_for_write(), mode="w") as f:
diff --git a/src/datachain/query/metrics.py b/src/datachain/query/metrics.py
index 065ad83..6143ab1 100644
--- a/src/datachain/query/metrics.py
+++ b/src/datachain/query/metrics.py
@@ -1,3 +1,4 @@
+import os
 from typing import Optional, Union
 
 metrics: dict[str, Union[str, int, float, bool, None]] = {}
@@ -13,6 +14,13 @@ def set(key: str, value: Union[str, int, float, bool, None]) -> None:  # noqa: P
         raise TypeError("Value must be a string, int, float or bool")
     metrics[key] = value
 
+    if job_id := os.getenv("DATACHAIN_JOB_ID"):
+        from datachain.data_storage.job import JobStatus
+        from datachain.query.session import Session
+
+        metastore = Session.get().catalog.metastore
+        metastore.set_job_status(job_id, JobStatus.RUNNING, metrics=metrics)
+
 
 def get(key: str) -> Optional[Union[str, int, float, bool]]:
     """Get a metric value."""


Test patch:
diff --git a/tests/func/test_metrics.py b/tests/func/test_metrics.py
new file mode 100644
index 0000000..57faa0c
--- /dev/null
+++ b/tests/func/test_metrics.py
@@ -0,0 +1,14 @@
+from datachain.query import metrics
+
+
+def test_metrics(test_session, catalog, mocker, monkeypatch):
+    mocker.patch("datachain.query.session.Session.get", return_value=test_session)
+    mocker.patch.dict("datachain.query.metrics.metrics", {}, clear=True)
+    job_id = catalog.metastore.create_job("test_metrics", "")
+    monkeypatch.setenv("DATACHAIN_JOB_ID", job_id)
+
+    metrics.set("foo", 42)
+
+    assert metrics.get("foo") == 42
+    (job,) = catalog.metastore.list_jobs_by_ids([job_id])
+    assert job.metrics == metrics.metrics == {"foo": 42}