|
14 | 14 | from pathlib import Path |
15 | 15 | from typing import Any |
16 | 16 |
|
17 | | -from braintrust.logger import Attachment, flush, init, init_dataset, init_logger |
| 17 | +from braintrust.logger import Attachment, flush, init, init_dataset |
18 | 18 | from exceptiongroup import ExceptionGroup |
19 | 19 |
|
20 | 20 | from .atif import _INSTRUMENTATION, ATIFImportResult, import_trajectory, summarize_trajectory |
@@ -263,7 +263,6 @@ def __init__( |
263 | 263 | max_attachment_bytes: Any = _UNSET, |
264 | 264 | max_total_attachment_bytes: Any = _UNSET, |
265 | 265 | max_content_bytes: Any = _UNSET, |
266 | | - log_job_summary: Any = _UNSET, |
267 | 266 | log_retry_attempts: Any = _UNSET, |
268 | 267 | strict: Any = _UNSET, |
269 | 268 | **kwargs: Any, |
@@ -291,7 +290,6 @@ def __init__( |
291 | 290 | "max_attachment_bytes": max_attachment_bytes, |
292 | 291 | "max_total_attachment_bytes": max_total_attachment_bytes, |
293 | 292 | "max_content_bytes": max_content_bytes, |
294 | | - "log_job_summary": log_job_summary, |
295 | 293 | "log_retry_attempts": log_retry_attempts, |
296 | 294 | "strict": strict, |
297 | 295 | **kwargs, |
@@ -356,7 +354,7 @@ async def reconcile(result: Any) -> None: |
356 | 354 | for name in set(self._trial_machines) - final_names: |
357 | 355 | await self._dispatch(name, TrialEvent(TrialEventKind.OMIT)) |
358 | 356 | try: |
359 | | - await asyncio.to_thread(self._finalize, job_result) |
| 357 | + await asyncio.to_thread(flush) |
360 | 358 | except Exception as exc: |
361 | 359 | failures.append(exc) |
362 | 360 | self._errors.append(f"final flush: {exc}") |
@@ -813,39 +811,6 @@ def _sync_final_result(self, result: Any) -> None: |
813 | 811 | root.log(metadata=metadata) |
814 | 812 | root.end(end_time=root_end) |
815 | 813 |
|
816 | | - def _finalize(self, job_result: Any) -> None: |
817 | | - if self._runtime is None: |
818 | | - return |
819 | | - if self.config.log_job_summary: |
820 | | - project_logger = init_logger( |
821 | | - project=self.config.project_name, |
822 | | - project_id=self.config.project_id, |
823 | | - set_current=False, |
824 | | - ) |
825 | | - now = datetime.now().timestamp() |
826 | | - summary_start, summary_end = _timing(job_result, now, now) |
827 | | - summary = project_logger.start_span( |
828 | | - name="harbor.job.summary", |
829 | | - type="task", |
830 | | - id=f"harbor-job-summary-{self._runtime.snapshot.job_id}", |
831 | | - start_time=summary_start, |
832 | | - set_current=False, |
833 | | - input={"job_id": self._runtime.snapshot.job_id}, |
834 | | - metadata={ |
835 | | - "harbor": { |
836 | | - "job_id": self._runtime.snapshot.job_id, |
837 | | - "experiments": [ |
838 | | - {"id": partition.experiment_id, "name": partition.name} |
839 | | - for partition in self._runtime.partitions.values() |
840 | | - ], |
841 | | - } |
842 | | - }, |
843 | | - internal={"instrumentation": _INSTRUMENTATION}, |
844 | | - ) |
845 | | - summary.log(output=job_result.stats.model_dump(mode="json", exclude_none=False)) |
846 | | - summary.end(end_time=summary_end) |
847 | | - flush() |
848 | | - |
849 | 814 | def _persist_disabled_manifest(self) -> None: |
850 | 815 | if self._snapshot is None: |
851 | 816 | return |
|
0 commit comments