From 746f419ceca38e3162417e2042d15af570513fa6 Mon Sep 17 00:00:00 2001 From: Lukasz Kolodziejczyk Date: Fri, 27 Jun 2025 13:47:01 +0200 Subject: [PATCH 1/4] wip --- mostlyai/sdk/_local/execution/jobs.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/mostlyai/sdk/_local/execution/jobs.py b/mostlyai/sdk/_local/execution/jobs.py index f374ced6..1aca9bc4 100644 --- a/mostlyai/sdk/_local/execution/jobs.py +++ b/mostlyai/sdk/_local/execution/jobs.py @@ -122,20 +122,25 @@ def _move_generation_artefacts(synthetic_dataset_dir: Path, job_workspace_dir: P def _mark_in_progress(resource: Generator | SyntheticDataset, resource_dir: Path): if isinstance(resource, Generator): resource.training_status = ProgressStatus.in_progress + write_resource_to_json = write_generator_to_json else: resource.generation_status = ProgressStatus.in_progress + write_resource_to_json = write_synthetic_dataset_to_json job_progress = read_job_progress_from_json(resource_dir) job_progress.status = ProgressStatus.in_progress for step in job_progress.steps: step.status = ProgressStatus.in_progress + write_resource_to_json(resource_dir, resource) write_job_progress_to_json(resource_dir, job_progress) def _mark_done(resource: Generator | SyntheticDataset, resource_dir: Path): if isinstance(resource, Generator): resource.training_status = ProgressStatus.done + write_resource_to_json = write_generator_to_json else: resource.generation_status = ProgressStatus.done + write_resource_to_json = write_synthetic_dataset_to_json now = get_current_utc_time() job_progress = read_job_progress_from_json(resource_dir) job_progress.status = ProgressStatus.done @@ -147,19 +152,23 @@ def _mark_done(resource: Generator | SyntheticDataset, resource_dir: Path): step.start_date = step.start_date or now step.end_date = step.end_date or now step.progress.value = step.progress.max + write_resource_to_json(resource_dir, resource) write_job_progress_to_json(resource_dir, job_progress) def _mark_failed(resource: Generator | SyntheticDataset, resource_dir: Path): if isinstance(resource, Generator): resource.training_status = ProgressStatus.failed + write_resource_to_json = write_generator_to_json else: resource.generation_status = ProgressStatus.failed + write_resource_to_json = write_synthetic_dataset_to_json job_progress = read_job_progress_from_json(resource_dir) job_progress.status = ProgressStatus.failed for step in job_progress.steps: if step.status != ProgressStatus.done: step.status = ProgressStatus.failed + write_resource_to_json(resource_dir, resource) write_job_progress_to_json(resource_dir, job_progress) @@ -607,7 +616,6 @@ def execute_training_job(generator_id: str, home_dir: Path): finally: execution.clear_job_workspace() execution.clear_file_upload_connectors() - write_generator_to_json(generator_dir, generator) try: _probe_random_samples(home_dir=home_dir, generator=generator) @@ -652,7 +660,6 @@ def execute_generation_job(synthetic_dataset_id: str, home_dir: Path): finally: execution.clear_job_workspace() execution.clear_file_upload_connectors() - write_synthetic_dataset_to_json(synthetic_dataset_dir, synthetic_dataset) def execute_probing_job(synthetic_dataset_id: str, home_dir: Path) -> list[Probe]: From 83b1b32722952ca78f1b2d705bf57e3de9ecfb17 Mon Sep 17 00:00:00 2001 From: Lukasz Kolodziejczyk Date: Fri, 27 Jun 2025 14:08:40 +0200 Subject: [PATCH 2/4] wip --- mostlyai/sdk/_local/progress.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/mostlyai/sdk/_local/progress.py b/mostlyai/sdk/_local/progress.py index 2f3016fd..a627a1fa 100644 --- a/mostlyai/sdk/_local/progress.py +++ b/mostlyai/sdk/_local/progress.py @@ -13,14 +13,13 @@ # limitations under the License. import datetime -import json import math from pathlib import Path from pydantic import BaseModel -from mostlyai.sdk._local.storage import write_to_json -from mostlyai.sdk.domain import JobProgress, ProgressStatus, StepCode +from mostlyai.sdk._local.storage import read_job_progress_from_json +from mostlyai.sdk.domain import ProgressStatus, StepCode def get_current_utc_time() -> datetime.datetime: @@ -43,8 +42,7 @@ def __init__(self, resource_path: Path, model_label: str | None, step_code: Step self._last_send_progress_time = None self._total = None - self.progress_file = self.resource_path / "job_progress.json" - self.job_progress = JobProgress(**json.loads(self.progress_file.read_text())) + self.job_progress = read_job_progress_from_json(self.resource_path) def _check_elapsed_interval(self): now = get_current_utc_time() @@ -133,6 +131,8 @@ def __call__( or (completed is not None and elapsed_enough_time) or (increase_by > 0 and elapsed_enough_time) ): - write_to_json(self.progress_file, self.job_progress) + # TODO: either update both resource and job_progress, or none of them + # write_job_progress_to_json(self.resource_path, self.job_progress) + pass return {} From 97920ab9a246cb8d170607df3a91c4cc41914e5d Mon Sep 17 00:00:00 2001 From: Lukasz Kolodziejczyk Date: Fri, 27 Jun 2025 14:54:14 +0200 Subject: [PATCH 3/4] wip --- mostlyai/sdk/_local/progress.py | 11 ++++------- mostlyai/sdk/client/_utils.py | 22 ++++++++++------------ 2 files changed, 14 insertions(+), 19 deletions(-) diff --git a/mostlyai/sdk/_local/progress.py b/mostlyai/sdk/_local/progress.py index a627a1fa..58e3c189 100644 --- a/mostlyai/sdk/_local/progress.py +++ b/mostlyai/sdk/_local/progress.py @@ -18,7 +18,7 @@ from pydantic import BaseModel -from mostlyai.sdk._local.storage import read_job_progress_from_json +from mostlyai.sdk._local.storage import read_job_progress_from_json, write_job_progress_to_json from mostlyai.sdk.domain import ProgressStatus, StepCode @@ -117,9 +117,8 @@ def __call__( self.job_progress.progress.max = len(self.job_progress.steps) # of steps if self.job_progress.start_date is None: self.job_progress.start_date = now - if self.job_progress.progress.value >= self.job_progress.progress.max: - self.job_progress.end_date = now - self.job_progress.status = ProgressStatus.done + # NOTE: do not set job status to DONE here, that should happen as a very last thing + # in order for _job_wait not to finish polling prematurely # send progress if we are DONE, or if we have a message to pass, # or if enough time has passed since last progress update @@ -131,8 +130,6 @@ def __call__( or (completed is not None and elapsed_enough_time) or (increase_by > 0 and elapsed_enough_time) ): - # TODO: either update both resource and job_progress, or none of them - # write_job_progress_to_json(self.resource_path, self.job_progress) - pass + write_job_progress_to_json(self.resource_path, self.job_progress) return {} diff --git a/mostlyai/sdk/client/_utils.py b/mostlyai/sdk/client/_utils.py index a020f0e9..3b2f04b4 100644 --- a/mostlyai/sdk/client/_utils.py +++ b/mostlyai/sdk/client/_utils.py @@ -288,18 +288,16 @@ def _rich_table_insert_row(*renderables: RenderableType | None, table: Table, id if step.status in (ProgressStatus.failed, ProgressStatus.canceled): rich.print(f"[red]Step {step.model_label} {step.step_code.value} {step.status.lower()}") return - # check whether we are done - if job.progress.value >= job.progress.max: - live.refresh() - time.sleep(1) # give the system a moment to update the status - return - else: - if job.end_date or job.progress in ( - ProgressStatus.failed, - ProgressStatus.canceled, - ): - rich.print(f"Job {job.status.lower()}") - return + + # check whether we are done + if job.status in ( + ProgressStatus.done, + ProgressStatus.failed, + ProgressStatus.canceled, + ): + live.refresh() + time.sleep(1) # give the system a moment to update the status + return except KeyboardInterrupt: rich.print(f"[red]Step {step.model_label} {step.step_code.value} {step.status.lower()}") return From bd5d1f1e2adb2fe83ed7d9b37624767aa1dc73e6 Mon Sep 17 00:00:00 2001 From: Lukasz Kolodziejczyk Date: Fri, 27 Jun 2025 16:15:32 +0200 Subject: [PATCH 4/4] wip --- mostlyai/sdk/_local/progress.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mostlyai/sdk/_local/progress.py b/mostlyai/sdk/_local/progress.py index 58e3c189..eb8f784e 100644 --- a/mostlyai/sdk/_local/progress.py +++ b/mostlyai/sdk/_local/progress.py @@ -117,8 +117,8 @@ def __call__( self.job_progress.progress.max = len(self.job_progress.steps) # of steps if self.job_progress.start_date is None: self.job_progress.start_date = now - # NOTE: do not set job status to DONE here, that should happen as a very last thing - # in order for _job_wait not to finish polling prematurely + # NOTE: do not set job status to DONE here, that should happen as the very last thing + # in order for `job_wait` not to finish polling prematurely # send progress if we are DONE, or if we have a message to pass, # or if enough time has passed since last progress update