From 0640da9e5db16e5d19614aae100a1ef196ac4152 Mon Sep 17 00:00:00 2001 From: mohamedelabbas1996 Date: Thu, 8 Oct 2026 16:23:42 -0400 Subject: [PATCH 1/8] feat(jobs): retrain a classifier head and take the result back Adds the training job: it builds the dataset, hands a processing service the URL, and records what comes back as a new algorithm version. Training outlasts the request that starts it, so the service reports its result to a callback instead of the connection being held open. The service has no Antenna account, so the callbacks are authorised by a signed token issued when the job was dispatched, and the job is looked up directly because project visibility would refuse an unauthenticated caller. The weights are uploaded the same way: without that, the head exists only in a cache directory on the service's disk and Antenna records a version it cannot point at. An upload is refused once the job has finished, since the token lasts a day and the path is fixed by the job, so otherwise the weights behind a registered version could be swapped out after the run ended. A job type now also declares what it needs before it can start and whether a user may create it at all, rather than each job discovering a missing param partway through its own run. --- .../0024_train_classifier_job_type.py | 31 ++ ami/jobs/models.py | 369 +++++++++++++++++- ami/jobs/views.py | 106 +++++ ami/ml/test_training.py | 39 ++ ami/ml/training/__init__.py | 32 +- ami/ml/training/service.py | 199 ++++++++++ 6 files changed, 772 insertions(+), 4 deletions(-) create mode 100644 ami/jobs/migrations/0024_train_classifier_job_type.py create mode 100644 ami/ml/training/service.py diff --git a/ami/jobs/migrations/0024_train_classifier_job_type.py b/ami/jobs/migrations/0024_train_classifier_job_type.py new file mode 100644 index 000000000..314738095 --- /dev/null +++ b/ami/jobs/migrations/0024_train_classifier_job_type.py @@ -0,0 +1,31 @@ +# Generated by Django 4.2.10 on 2026-10-08 15:53 + +from django.db import migrations, models + + +class Migration(migrations.Migration): + dependencies = [ + ("jobs", "0023_alter_job_job_type_key"), + ] + + operations = [ + migrations.AlterField( + model_name="job", + name="job_type_key", + field=models.CharField( + choices=[ + ("ml", "ML pipeline"), + ("populate_captures_collection", "Populate capture set"), + ("data_storage_sync", "Data storage sync"), + ("regroup_events", "Regroup sessions"), + ("unknown", "Unknown"), + ("data_export", "Data Export"), + ("post_processing", "Post Processing"), + ("train_classifier", "Train classifier"), + ], + default="unknown", + max_length=255, + verbose_name="Job Type", + ), + ), + ] diff --git a/ami/jobs/models.py b/ami/jobs/models.py index 20d03a54c..4d0607d6c 100644 --- a/ami/jobs/models.py +++ b/ami/jobs/models.py @@ -997,6 +997,371 @@ def run(cls, job: "Job"): raise ValueError(f"Unknown job type '{job.job_type()}'") +class TrainClassifierJob(JobType): + """ + Retrain a classifier head from the species people have verified in this project. + + Antenna prepares the training set and hands the processing service a URL to it, rather + than posting the rows: a project with a few hundred thousand verified labels runs to + hundreds of megabytes, which is fragile to send in one request and has to start over if + the connection drops. + + The job does not wait for training to finish. It dispatches and leaves the job STARTED; + the service reports back through the job's result endpoint. Waiting inline would hit the + sub-task timeout on anything but a toy dataset. + """ + + name = "Train classifier" + key = "train_classifier" + required_params = ("algorithm_key", "occurrence_set_id") + user_creatable = True + + # This job dispatches and then waits for the service to report back, writing nothing to + # its own row in between, so the default threshold reads a healthy run as a dead one. + # Matched to the life of the callback token: once that expires Antenna refuses the + # callback anyway, so a job still waiting past it can never finish. + stalled_after_minutes = 24 * 60 + + STAGE_PREPARE = "prepare" + STAGE_DISPATCH = "dispatch" + STAGE_TRAIN = "train" + + @classmethod + def run(cls, job: "Job"): + from ami.ml.models import Algorithm + from ami.ml.models.processing_service import ProcessingService + from ami.ml.training_dataset import NotEnoughVerifiedData, build_training_dataset + + params = job.params or {} + algorithm_key = params.get("algorithm_key") + if not algorithm_key: + raise ValueError("A train_classifier job needs an 'algorithm_key' in its params.") + + algorithm = Algorithm.objects.filter(key=algorithm_key).first() + if not algorithm: + raise ValueError(f"No algorithm with key '{algorithm_key}'.") + if not algorithm.trainable: + raise ValueError( + f"Algorithm '{algorithm_key}' is not marked trainable by its processing service. " + "Re-register the pipelines if the service has since been updated." + ) + + # Settled before any stage exists, so a bad setting is refused outright rather than + # failing halfway through a run that already looks started. + occurrence_set = cls.target_occurrence_set(job) + config = cls.config_for(job, algorithm) + + job.progress.add_stage("Preparing training set", cls.STAGE_PREPARE) + job.progress.add_stage("Sending to processing service", cls.STAGE_DISPATCH) + job.progress.add_stage("Training", cls.STAGE_TRAIN) + job.update_status(JobState.STARTED) + job.started_at = datetime.datetime.now() + job.finished_at = None + job.progress.update_stage(cls.STAGE_PREPARE, status=JobState.STARTED, progress=0) + job.save() + + taxa_list = cls.target_taxa_list(job) + try: + dataset = build_training_dataset( + project=job.project, + algorithm=algorithm, + min_per_species=config.min_per_species, + test_fraction=config.test_fraction, + split_salt=config.split_salt, + taxa_list=taxa_list, + occurrence_set=occurrence_set, + job=job, + ) + except NotEnoughVerifiedData as e: + # A data problem, not a crash. Say so plainly rather than failing with a traceback. + job.logger.error(str(e)) + job.progress.update_stage(cls.STAGE_PREPARE, status=JobState.FAILURE, progress=0) + job.finished_at = datetime.datetime.now() + job.result = {"error": str(e)} + job.update_status(JobState.FAILURE, save=True) + return + + meta = dataset["metadata"] + job.logger.info( + f"Training set: {meta.rows} verified crops over {len(meta.classes)} species " + f"({meta.train} train / {meta.test} held out)" + ) + if meta.taxa_list: + job.logger.info( + f"Species list comes from taxa list '{meta.taxa_list.name}' " f"({len(meta.classes)} species)" + ) + without_data = meta.classes_without_verified_data or [] + if without_data: + job.logger.info( + f"{len(without_data)} species in the list have no verified crops yet. They keep " + "whatever the current head already knows." + ) + if meta.dropped_species: + job.logger.warning( + f"{len(meta.dropped_species)} verified species are not in the taxa list and were " + f"left out: {', '.join(meta.dropped_species[:5])}" + f"{' ...' if len(meta.dropped_species) > 5 else ''}" + ) + if meta.verified_detections_without_embedding: + job.logger.warning( + f"{meta.verified_detections_without_embedding} verified detection(s) have no embedding " + "from this algorithm and were left out. Re-run the pipeline over them to include them." + ) + job.progress.add_stage_param(cls.STAGE_PREPARE, "Rows", meta.rows) + job.progress.add_stage_param(cls.STAGE_PREPARE, "Species", len(meta.classes)) + job.progress.add_stage_param(cls.STAGE_PREPARE, "Dataset", dataset["url"]) + job.progress.update_stage(cls.STAGE_PREPARE, status=JobState.SUCCESS, progress=1) + job.save() + + service = ( + ProcessingService.objects.filter( + projects=job.project, + pipelines__algorithms=algorithm, + endpoint_url__isnull=False, + ) + .exclude(endpoint_url="") + .distinct() + .first() + ) + if not service: + # Pull-mode workers register with a null endpoint_url, so there is nothing to + # send a training request to. Supporting them means routing this through the + # task queue instead. + raise ValueError( + f"No push-mode processing service in this project serves '{algorithm_key}'. " + "Training cannot be dispatched to a pull-mode worker." + ) + + job.progress.update_stage(cls.STAGE_DISPATCH, status=JobState.STARTED, progress=0) + job.save() + cls.dispatch(job=job, service=service, algorithm=algorithm, dataset=dataset) + + @classmethod + def config_for(cls, job: "Job", algorithm): + """ + The training settings for this run, checked before anything uses them. + + The algorithm's training_config holds the defaults its processing service + published; a job may override any of them for one run. Those overrides arrive in + job.params, which any member who can create a job sets, so they are validated here + instead of being read raw at each point of use. + """ + import pydantic + + from ami.ml.schemas import AlgorithmTrainingConfig + + try: + return AlgorithmTrainingConfig.for_run(algorithm.training_config, job.params or {}) + except pydantic.ValidationError as e: + problems = "; ".join(f"{'.'.join(str(part) for part in err['loc'])} {err['msg']}" for err in e.errors()) + raise ValueError(f"This job's training settings are out of range: {problems}") + + @classmethod + def target_occurrence_set(cls, job: "Job"): + """ + The set of occurrences this run learns from. + + Required rather than optional: a set is fixed once created, so a run that names one + can be repeated and says for itself what it learned from. Training on whatever + happened to be verified that day leaves nothing a later reader can reconstruct. + """ + from ami.main.models import OccurrenceSet + + occurrence_set_id = (job.params or {}).get("occurrence_set_id") + if not occurrence_set_id: + raise ValueError("A train_classifier job needs an 'occurrence_set_id' in its params.") + # Scoped to the project, since params are set by any member who can create a job. + occurrence_set = OccurrenceSet.objects.for_project(job.project).filter(pk=occurrence_set_id).first() + if not occurrence_set: + raise ValueError(f"No occurrence set with id {occurrence_set_id} in this project.") + return occurrence_set + + @classmethod + def target_taxa_list(cls, job: "Job"): + """ + Which taxa list sets the head's species. + + A job may name one; otherwise the project's default is used. Without either, the + class list falls back to whatever has been verified. + """ + from ami.main.models import TaxaList + + taxa_list_id = (job.params or {}).get("taxa_list_id") + if taxa_list_id: + # Scoped for the same reason as the evaluation set above. + taxa_list = TaxaList.objects.for_project(job.project).filter(pk=taxa_list_id).first() + if not taxa_list: + raise ValueError(f"No taxa list with id {taxa_list_id} in this project.") + return taxa_list + return job.project.default_taxa_list + + @classmethod + def dispatch(cls, job: "Job", service, algorithm, dataset: dict) -> None: + """Hand the service the dataset URL and leave the job running until it reports back.""" + from ami.ml.training_dispatch import send_training_request + + response = send_training_request(job=job, service=service, algorithm=algorithm, dataset=dataset) + + # The service posts its callback before returning from /train, so a fast run is + # already finished by the time this line is reached. Writing this copy's progress + # would put the training stage back to "started" on a job that is done. + job.refresh_from_db() + if job.status in JobState.final_states(): + job.logger.info(f"{service.name} already reported its result.") + return + + job.progress.update_stage(cls.STAGE_DISPATCH, status=JobState.SUCCESS, progress=1) + job.progress.update_stage(cls.STAGE_TRAIN, status=JobState.STARTED, progress=0) + job.logger.info(f"Training request accepted by {service.name}. Waiting for it to report back.") + job.save(update_fields=["progress", "updated_at"]) + + if response is None: + # The service accepted the work and will post to the job's training-result + # endpoint when it finishes. Leaving the job STARTED is the point: a real + # training set outlasts the request that started it. + job.logger.info("Waiting for the service to report back.") + return + + if response is not None: + # Answered inline, which small datasets do. Record it now rather than waiting + # for a callback that has already been overtaken. + cls.record_result( + job=job, + payload={ + "result": response, + "dataset": dataset["metadata"], + # Kept alongside the metadata so the registered version can point at the + # exact file it was fitted on, not just describe it. + "dataset_url": dataset["url"], + }, + ) + + @classmethod + def record_result(cls, job: "Job", payload: dict) -> None: + """Store what the service reported, register the new version, and finish the job.""" + # A service posts its callback before returning from /train, so a fast run reports + # twice: once through the callback and once inline. Without this guard each retrain + # registered two algorithm versions. Re-read first, because the inline caller holds + # a copy from before the callback landed. + job.refresh_from_db(fields=["status", "result"]) + if job.status in JobState.final_states(): + job.logger.info("A training result is already recorded for this job; ignoring a duplicate.") + return + + result = payload.get("result") or {} + job.result = payload + + # A result can arrive through the callback on a job whose stages were never set up, + # for instance after a restart. Make sure the stage exists before reporting into it. + if not any(stage.key == cls.STAGE_TRAIN for stage in job.progress.stages): + job.progress.add_stage("Training", cls.STAGE_TRAIN) + # A service that is reporting back plainly received the request, and the callback + # can arrive before the dispatching code closes that stage. + if any(stage.key == cls.STAGE_DISPATCH for stage in job.progress.stages): + job.progress.update_stage(cls.STAGE_DISPATCH, status=JobState.SUCCESS, progress=1) + + for warning in result.get("warnings", []): + job.logger.warning(warning) + + candidate = result.get("candidate_metrics") or {} + incumbent = result.get("incumbent_metrics") or {} + job.progress.add_stage_param(cls.STAGE_TRAIN, "New head top-1", candidate.get("top1")) + job.progress.add_stage_param(cls.STAGE_TRAIN, "Current head top-1", incumbent.get("top1")) + job.progress.add_stage_param(cls.STAGE_TRAIN, "Better", result.get("promote")) + job.logger.info(result.get("reason", "Training finished.")) + + new_version = cls.register_new_version(job=job, payload=payload) + if new_version: + job.progress.add_stage_param(cls.STAGE_TRAIN, "New version", new_version.key) + job.logger.info(f"Registered {new_version} as version {new_version.version}") + + job.progress.update_stage(cls.STAGE_TRAIN, status=JobState.SUCCESS, progress=1) + job.finished_at = datetime.datetime.now() + job.update_status(JobState.SUCCESS, save=True) + job.save() + + @classmethod + def register_new_version(cls, job: "Job", payload: dict): + """ + Record the trained head as a new version of the algorithm it was trained from. + + A Classification points at an Algorithm row, so a version that changed in place + would make past predictions untraceable. Each retrain therefore gets its own row: + same name, next version number, and a new key, because key is unique on its own. + """ + import pydantic + + from ami.ml.models import Algorithm, AlgorithmCategoryMap + from ami.ml.schemas import AlgorithmTrainingInfo, TrainingDatasetMetadata + + result = payload.get("result") or {} + dataset = payload.get("dataset") or {} + # The service echoes the dataset's metadata back as JSON, so it arrives as a dict + # rather than the model that wrote it. Parsed here so a malformed echo is caught at + # the boundary instead of further in. + metadata = None + if dataset.get("metadata"): + try: + metadata = TrainingDatasetMetadata.parse_obj(dataset["metadata"]) + except pydantic.ValidationError as e: + job.logger.warning(f"The service returned training set metadata Antenna could not read: {e}") + parent_key = (job.params or {}).get("algorithm_key") + parent = Algorithm.objects.filter(key=parent_key).first() + if not parent: + job.logger.warning(f"Cannot register a new version: no algorithm with key '{parent_key}'.") + return None + + labels = result.get("labels") or [] + if not labels: + job.logger.warning("The service returned no class list, so no new version was registered.") + return None + + trained_at = result.get("trained_at") or "" + stamp = str(trained_at).replace(":", "").replace("-", "").replace(".", "")[:15] or f"job{job.pk}" + version = ( + Algorithm.objects.filter(name=parent.name).order_by("-version").values_list("version", flat=True).first() + or parent.version + ) + 1 + + category_map = AlgorithmCategoryMap.objects.create( + labels=labels, + data=[{"label": label, "index": i} for i, label in enumerate(labels)], + version=stamp, + description=( + f"Retrained from {dataset.get('rows', len(labels))} verified crops in " + f"project '{job.project.name}' by job #{job.pk}." + ), + ) + + return Algorithm.objects.create( + name=parent.name, + key=f"{parent.key}-v{version}-{stamp}", + # Where the weights are kept. Empty when the service did not upload them, in + # which case the head exists only on that service's own disk. + uri=result.get("head_url") or "", + version=version, + version_name=str(trained_at) or stamp, + task_type=parent.task_type, + description=parent.description, + trainable=parent.trainable, + training_config=parent.training_config, + category_map=category_map, + training_info=AlgorithmTrainingInfo( + trained_at=trained_at or None, + dataset_url=payload.get("dataset_url") or dataset.get("url"), + occurrence_set_id=metadata.occurrence_set.id if metadata and metadata.occurrence_set else None, + occurrence_set_name=metadata.occurrence_set.name if metadata and metadata.occurrence_set else None, + dataset_rows=(result.get("rows") or {}).get("kept"), + dataset_classes=len(labels), + metrics=result.get("candidate_metrics") or {}, + previous_metrics=result.get("incumbent_metrics") or {}, + parent_algorithm_key=parent.key, + job_id=job.pk, + warnings=result.get("warnings") or [], + ), + ) + + VALID_JOB_TYPES = [ MLJob, SourceImageCollectionPopulateJob, @@ -1005,6 +1370,7 @@ def run(cls, job: "Job"): UnknownJobType, DataExportJob, PostProcessingJob, + TrainClassifierJob, ] @@ -1041,7 +1407,8 @@ class Job(BaseModel): # Redis SREM-driven progress save, so this is effectively "no progress for # N minutes". 10 is conservative; raise if legitimate long-running jobs get # reaped. - STALLED_JOBS_MAX_MINUTES = 10 + # The default deadline, owned by JobType so a job type can set its own. + STALLED_JOBS_MAX_MINUTES = JobType.stalled_after_minutes # Zombie-stream reaper: age threshold above which a NATS stream for a job # in a terminal state (or missing from Django) is considered safe to drop. # Kept well above :attr:`STALLED_JOBS_MAX_MINUTES` so newly-dispatched jobs diff --git a/ami/jobs/views.py b/ami/jobs/views.py index 09b314e4b..62cc64dc9 100644 --- a/ami/jobs/views.py +++ b/ami/jobs/views.py @@ -7,6 +7,7 @@ from django.core.cache import cache from django.db.models import Q from django.db.models.query import QuerySet +from django.http import Http404 from django.utils import timezone from django_filters import rest_framework as filters from drf_spectacular.utils import extend_schema, extend_schema_view @@ -14,6 +15,8 @@ from rest_framework.decorators import action from rest_framework.exceptions import NotFound, PermissionDenied, ValidationError from rest_framework.filters import BaseFilterBackend +from rest_framework.parsers import MultiPartParser +from rest_framework.permissions import AllowAny from rest_framework.response import Response from ami.base.filters import RelatedIdFilter @@ -591,3 +594,106 @@ def result(self, request, pk=None): }, status=503, ) + + def _job_for_callback(self, pk, request) -> Job: + """ + The job a processing service is reporting about, if its token proves it may. + + get_object() applies project visibility, which an unauthenticated service fails, + so the job is looked up directly and the signed token is what authorises the call. + """ + from ami.ml.training import verify_callback_token + + job = Job.objects.filter(pk=pk).first() + if not job: + raise Http404("Job not found.") + + token = request.headers.get("Authorization", "").removeprefix("Token ").strip() + if not verify_callback_token(token, job): + raise PermissionDenied("Invalid or expired training callback token.") + return job + + @extend_schema( + request=MLJobResultsRequestSerializer, + responses={200: MLJobResultsResponseSerializer}, + parameters=[project_id_doc_param], + ) + @extend_schema(exclude=True) + @action( + detail=True, + methods=["post"], + url_path="training-result", + name="training-result", + # A processing service has no Antenna account. It proves itself with the signed + # token Antenna issued when it dispatched the job, checked below. + permission_classes=[AllowAny], + authentication_classes=[], + ) + def training_result(self, request, pk=None): + """ + Receive the outcome of a retraining run from a processing service. + + Training can outlast the request that started it, so the service reports back here + instead of holding the connection open. + """ + from ami.jobs.models import TrainClassifierJob + + job = self._job_for_callback(pk, request) + + if job.job_type_key != TrainClassifierJob.key: + raise ValidationError(f"Job #{job.pk} is not a training job.") + + if job.status in JobState.final_states(): + # The service answered inline and the result is already recorded, or a retry + # arrived late. Either way the first answer stands. + logger.info("Ignoring a training result for job %s, which already finished", job.pk) + return Response({"status": "already recorded"}) + + TrainClassifierJob.record_result(job=job, payload=request.data) + logger.info("Recorded a training result for job %s", job.pk) + return Response({"status": "recorded"}) + + @extend_schema(exclude=True) + @action( + detail=True, + methods=["post"], + url_path="training-head", + name="training-head", + # Same token as the result callback: a processing service has no Antenna account. + permission_classes=[AllowAny], + authentication_classes=[], + parser_classes=[MultiPartParser], + ) + def training_head(self, request, pk=None): + """ + Receive the head a retraining run produced, so Antenna keeps a copy of the weights. + + Without this the head exists only on the service's disk, under a cache directory, + and Antenna records that a version exists without being able to say where it is. + """ + from ami.jobs.models import TrainClassifierJob + from ami.ml.training import HeadTooLarge, store_head + + job = self._job_for_callback(pk, request) + + if job.job_type_key != TrainClassifierJob.key: + raise ValidationError(f"Job #{job.pk} is not a training job.") + + if job.status in JobState.final_states(): + # The token stays valid for 24 hours and the head is stored at a path fixed by + # the job, so without this the weights a registered version points at could be + # replaced for a day after the run ended. A service uploads before it reports + # its result, so a run that is still going is unaffected. + raise ValidationError(f"Job #{job.pk} has already finished; its head can no longer be replaced.") + + if not request.FILES: + raise ValidationError("No head files were uploaded.") + + algorithm_key = (job.params or {}).get("algorithm_key") or "head" + try: + stored = store_head(algorithm_key=algorithm_key, job_id=job.pk, files=request.FILES) + except HeadTooLarge as e: + raise ValidationError(str(e)) + + job.logger.info(f"Stored the retrained head: {', '.join(item['path'] for item in stored.values())}") + return Response({"files": stored}) diff --git a/ami/ml/test_training.py b/ami/ml/test_training.py index 896e7fc7a..bf4c8fad9 100644 --- a/ami/ml/test_training.py +++ b/ami/ml/test_training.py @@ -258,3 +258,42 @@ def test_someone_who_cannot_retrain_cannot_read_the_training_set(self): response = self.client.get(self.url, {"project_id": self.project.pk, "algorithm": self.algorithm.key}) self.assertEqual(response.status_code, 403) + + +class TestTheJobRequiresASet(TrainingSetFixture): + def _job(self, **params): + from ami.jobs.models import Job, TrainClassifierJob + + return Job.objects.create( + project=self.project, + name="Retrain", + job_type_key=TrainClassifierJob.key, + params={"algorithm_key": self.algorithm.key, **params}, + ) + + def test_a_job_without_a_set_is_refused(self): + from ami.jobs.models import TrainClassifierJob + + with self.assertRaises(ValueError) as caught: + TrainClassifierJob.target_occurrence_set(self._job()) + + self.assertIn("occurrence_set_id", str(caught.exception)) + + def test_a_set_from_another_project_is_refused(self): + from ami.jobs.models import TrainClassifierJob + + other = Project.objects.create(name="Someone else", owner=self.user) + theirs = OccurrenceSet.objects.create(name="Theirs") + theirs.projects.add(other) + + with self.assertRaises(ValueError) as caught: + TrainClassifierJob.target_occurrence_set(self._job(occurrence_set_id=theirs.pk)) + + self.assertIn("in this project", str(caught.exception)) + + def test_a_set_in_this_project_is_used(self): + from ami.jobs.models import TrainClassifierJob + + mine = self.make_set("Mine", [self.make_occurrence(self.taxa[0])]) + + self.assertEqual(TrainClassifierJob.target_occurrence_set(self._job(occurrence_set_id=mine.pk)), mine) diff --git a/ami/ml/training/__init__.py b/ami/ml/training/__init__.py index 6770457ad..268c6c6d5 100644 --- a/ami/ml/training/__init__.py +++ b/ami/ml/training/__init__.py @@ -1,9 +1,9 @@ """ Retraining a classifier head from the species people have verified. -``dataset`` decides what the head learns from and writes it to storage. The names -re-exported here are the ones callers outside this package use; anything else is internal -to it. +``dataset`` decides what the head learns from and writes it to storage; ``service`` hands a +processing service the URL and keeps the head that comes back. The names re-exported here +are the ones callers outside this package use; anything else is internal to it. """ from ami.ml.training.dataset import ( @@ -22,20 +22,46 @@ verified_occurrence_ids, verified_training_rows, ) +from ami.ml.training.service import ( + CALLBACK_MAX_AGE_SECONDS, + DISPATCH_TIMEOUT_SECONDS, + MAX_HEAD_BYTES, + HeadTooLarge, + absolute_media_url, + callback_url_for, + head_path, + head_upload_url_for, + make_callback_token, + send_training_request, + store_head, + verify_callback_token, +) __all__ = [ + "CALLBACK_MAX_AGE_SECONDS", "DEFAULT_SPLIT_SALT", "DEFAULT_TEST_FRACTION", + "DISPATCH_TIMEOUT_SECONDS", + "MAX_HEAD_BYTES", "SPLITS", "SPLIT_TEST", "SPLIT_TRAIN", + "HeadTooLarge", "NotEnoughVerifiedData", + "absolute_media_url", "build_training_dataset", + "callback_url_for", "count_missing_embeddings", + "head_path", + "head_upload_url_for", "label_counts", + "make_callback_token", "row_as_dict", + "send_training_request", "species_with_enough_examples", "split_for", + "store_head", "verified_occurrence_ids", "verified_training_rows", + "verify_callback_token", ] diff --git a/ami/ml/training/service.py b/ami/ml/training/service.py new file mode 100644 index 000000000..b9a7c4920 --- /dev/null +++ b/ami/ml/training/service.py @@ -0,0 +1,199 @@ +""" +Hand a training request to a processing service, and keep what it sends back. + +Outward: a URL to the dataset rather than the rows (see ``dataset``), plus a signed token +the service reports back with. + +Inward: the service uploads the finished head through Antenna rather than to storage. It +holds no storage credentials, and a presigned upload would only work on S3, not against the +local filesystem. A head is one small matrix, so passing it through costs little. +""" + +import logging +import typing +from typing import TYPE_CHECKING +from urllib.parse import urljoin + +from django.conf import settings +from django.core import signing +from django.core.files.storage import default_storage +from django.urls import reverse +from django.utils.text import slugify + +from ami.ml.schemas import AlgorithmTrainingConfig +from ami.utils.requests import create_session, extract_error_message_from_response + +if TYPE_CHECKING: + from ami.jobs.models import Job + from ami.ml.models import Algorithm, ProcessingService + +logger = logging.getLogger(__name__) + + +# How long to wait for the service to acknowledge the request. This is not how long +# training takes: a service that answers within this window returns its result inline, +# and one that does not is expected to report back through the job's result endpoint. +DISPATCH_TIMEOUT_SECONDS = 600 + + +# A processing service has no Antenna account, so the callback is authorised by a signed +# token instead. Nothing is stored: the signature carries the job id and Django's secret +# key proves Antenna issued it. +CALLBACK_SALT = "ami.ml.training.callback" +CALLBACK_MAX_AGE_SECONDS = 60 * 60 * 24 + + +def make_callback_token(job: "Job") -> str: + """A token only Antenna could have produced, tied to this one job.""" + return signing.dumps({"job_id": job.pk}, salt=CALLBACK_SALT) + + +def verify_callback_token(token: str, job: "Job") -> bool: + """True when the token is Antenna's, unexpired, and for this job.""" + if not token: + return False + try: + payload = signing.loads(token, salt=CALLBACK_SALT, max_age=CALLBACK_MAX_AGE_SECONDS) + except signing.BadSignature: + return False + return payload.get("job_id") == job.pk + + +def _callback_base(job: "Job") -> str: + """ + Where the service can reach this Antenna. + + Deliberately not taken from the job's params. Params are free JSON that any member who + can create a job may set, and this base decides where the signed callback token is + sent -- a token that also unlocks the head upload. A per-deployment override belongs + in settings, not in a request body. + """ + base = getattr(settings, "EXTERNAL_BASE_URL", "") + if not base: + raise ValueError( + "No base URL is configured, so the processing service has no way to report back. " "Set EXTERNAL_BASE_URL." + ) + return base + + +def callback_url_for(job: "Job") -> str: + """Where the service should post its result when training finishes.""" + # reverse(), so the path follows the router rather than a copy of it here. + path = reverse("api:job-training-result", args=[job.pk]) + return urljoin(_callback_base(job).rstrip("/") + "/", path.lstrip("/")) + + +def head_upload_url_for(job: "Job") -> str: + """Where the service should upload the head it produced, so Antenna keeps a copy.""" + path = reverse("api:job-training-head", args=[job.pk]) + return urljoin(_callback_base(job).rstrip("/") + "/", path.lstrip("/")) + + +def absolute_media_url(url: str, base_url: str | None = None) -> str: + """ + Turn a stored file's URL into one a processing service can fetch. + + In production MEDIA_URL is already an absolute S3 URL and this is a no-op. Locally it is + the relative /media/ path, so it needs a base in front. + """ + if url.startswith("http://") or url.startswith("https://"): + return url + base = base_url or getattr(settings, "EXTERNAL_BASE_URL", "") + if not base: + raise ValueError( + "The training set is stored at a relative URL and no base URL is configured, so the " + "processing service has no way to download it. Set EXTERNAL_BASE_URL." + ) + return urljoin(base.rstrip("/") + "/", url.lstrip("/")) + + +def send_training_request( + job: "Job", service: "ProcessingService", algorithm: "Algorithm", dataset: dict +) -> dict | None: + """ + Ask a processing service to retrain a head. + + Returns the service's result if it answered inline, or None if it accepted the work and + will report back later. Raises if the service refused the request. + """ + endpoint = urljoin(service.endpoint_url.rstrip("/") + "/", "train") + # The same merge the job validated before it built the dataset, so the service is told + # the settings the training set was actually made under. Reading job.params again here + # would let the two drift, and would send values nothing had checked. + config = AlgorithmTrainingConfig.for_run(algorithm.training_config, job.params or {}) + payload: dict[str, typing.Any] = { + "dataset_url": absolute_media_url(dataset["url"]), + "algorithm_key": algorithm.key, + "job_id": job.pk, + "name": f"{algorithm.key}-job-{job.pk}", + "min_per_species": config.min_per_species, + # The fitting settings the service published, so an admin can tune them in Antenna + # without redeploying the service. + "min_improvement": config.min_improvement, + "head_type": config.head_type, + "epochs": config.epochs, + "learning_rate": config.learning_rate, + "weight_decay": config.weight_decay, + } + + # Always sent: a service that finishes after the request times out reports back here + # instead, which is the only way a real training set can work. + payload["callback_url"] = callback_url_for(job) + payload["callback_token"] = make_callback_token(job) + # The head itself comes back here. A service that does not support the upload simply + # ignores this, and the version is registered without a stored copy as before. + payload["head_upload_url"] = head_upload_url_for(job) + + job.logger.info(f"Sending training request to {endpoint} for {algorithm.key}") + session = create_session() + response = session.post(endpoint, json=payload, timeout=DISPATCH_TIMEOUT_SECONDS) + + if not response.ok: + message = extract_error_message_from_response(response) + raise ValueError(f"The processing service refused the training request: {message}") + + try: + return response.json() + except ValueError: + # Accepted, but nothing useful in the body. The service will report back. + return None + + +HEAD_DIRECTORY = "algorithms" + +# A head is one Linear layer: 72 KB over 17 species, about 3 MB over 749. The cap is well +# clear of that and only exists so a wrong or hostile upload cannot fill the bucket. +MAX_HEAD_BYTES = 64 * 1024 * 1024 + + +class HeadTooLarge(Exception): + """The upload is larger than any head this could reasonably be.""" + + +def head_path(algorithm_key: str, job_id: int, filename: str) -> str: + """Where one job's head is kept. Deterministic, so a re-run replaces its own file.""" + return f"{HEAD_DIRECTORY}/{slugify(algorithm_key)}-job-{job_id}/{filename}" + + +def store_head(algorithm_key: str, job_id: int, files: dict) -> dict[str, typing.Any]: + """ + Save the uploaded head files and return where they landed. + + Returns the storage path and URL of each file, keyed by the name it was uploaded + under, so the caller can record the head's location against the algorithm version. + """ + stored = {} + for name, uploaded in files.items(): + if uploaded.size > MAX_HEAD_BYTES: + raise HeadTooLarge(f"'{name}' is {uploaded.size} bytes; the limit is {MAX_HEAD_BYTES}.") + + path = head_path(algorithm_key, job_id, uploaded.name) + if default_storage.exists(path): + # A re-run of the same job replaces its head instead of piling up copies, + # matching how the training set is written. + default_storage.delete(path) + saved_path = default_storage.save(path, uploaded) + stored[name] = {"path": saved_path, "url": default_storage.url(saved_path)} + logger.info(f"Stored '{name}' for job {job_id} at {saved_path}") + + return stored From 4f253bf3f7d80aa025f97c47da0b3f24de7f66ae Mon Sep 17 00:00:00 2001 From: mohamedelabbas1996 Date: Thu, 8 Oct 2026 16:30:17 -0400 Subject: [PATCH 2/8] fix(jobs): point the training job at the merged training package The refactor that gathered the training modules into one package left two imports behind, so running a training job raised ImportError on ami.ml.training_dataset instead of building a dataset. The tests did not catch it because neither import is reached until a job actually runs. --- ami/jobs/models.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/ami/jobs/models.py b/ami/jobs/models.py index 4d0607d6c..e2ba50273 100644 --- a/ami/jobs/models.py +++ b/ami/jobs/models.py @@ -1030,7 +1030,7 @@ class TrainClassifierJob(JobType): def run(cls, job: "Job"): from ami.ml.models import Algorithm from ami.ml.models.processing_service import ProcessingService - from ami.ml.training_dataset import NotEnoughVerifiedData, build_training_dataset + from ami.ml.training import NotEnoughVerifiedData, build_training_dataset params = job.params or {} algorithm_key = params.get("algorithm_key") @@ -1198,7 +1198,7 @@ class list falls back to whatever has been verified. @classmethod def dispatch(cls, job: "Job", service, algorithm, dataset: dict) -> None: """Hand the service the dataset URL and leave the job running until it reports back.""" - from ami.ml.training_dispatch import send_training_request + from ami.ml.training import send_training_request response = send_training_request(job=job, service=service, algorithm=algorithm, dataset=dataset) From 6793f233eb2add98255af92d25d14e011aee71fb Mon Sep 17 00:00:00 2001 From: mohamedelabbas1996 Date: Thu, 8 Oct 2026 16:30:52 -0400 Subject: [PATCH 3/8] feat(jobs): make the training occurrence set optional Requiring a set meant that retraining on everything verified so far - the ordinary case - first made someone save a set of everything, which is a step with no decision in it. A job without a set now learns from every verified occurrence in the project. The dataset file records which set was used, or that there was none, so a run is still legible afterwards either way. Naming a set is still worth doing when it matters: a set is fixed once created, so that run can be repeated. --- ami/jobs/models.py | 22 ++++++++++++---------- ami/ml/test_training.py | 10 ++++------ 2 files changed, 16 insertions(+), 16 deletions(-) diff --git a/ami/jobs/models.py b/ami/jobs/models.py index e2ba50273..49505f81d 100644 --- a/ami/jobs/models.py +++ b/ami/jobs/models.py @@ -17,7 +17,7 @@ from ami.base.models import BaseModel from ami.base.schemas import ConfigurableStage, ConfigurableStageParam from ami.jobs.tasks import cleanup_async_job_if_needed, run_job -from ami.main.models import Deployment, Project, SourceImage, SourceImageCollection +from ami.main.models import Deployment, OccurrenceSet, Project, SourceImage, SourceImageCollection from ami.ml.models import Pipeline from ami.ml.post_processing.registry import get_postprocessing_task from ami.utils.schemas import OrderedEnum @@ -1013,7 +1013,7 @@ class TrainClassifierJob(JobType): name = "Train classifier" key = "train_classifier" - required_params = ("algorithm_key", "occurrence_set_id") + required_params = ("algorithm_key",) user_creatable = True # This job dispatches and then waits for the service to report back, writing nothing to @@ -1157,19 +1157,21 @@ def config_for(cls, job: "Job", algorithm): raise ValueError(f"This job's training settings are out of range: {problems}") @classmethod - def target_occurrence_set(cls, job: "Job"): + def target_occurrence_set(cls, job: "Job") -> OccurrenceSet | None: """ - The set of occurrences this run learns from. + The set of occurrences this run learns from, or None for every verified one. - Required rather than optional: a set is fixed once created, so a run that names one - can be repeated and says for itself what it learned from. Training on whatever - happened to be verified that day leaves nothing a later reader can reconstruct. - """ - from ami.main.models import OccurrenceSet + Naming a set is worth doing: a set is fixed once created, so the run can be repeated + and says for itself what it learned from. It is not required, because the common case + is retraining on everything verified so far, and making someone save a set first to + do that is a step with no decision in it. + The dataset file records which set was used, or that there was none, so a run without + one is still legible afterwards. + """ occurrence_set_id = (job.params or {}).get("occurrence_set_id") if not occurrence_set_id: - raise ValueError("A train_classifier job needs an 'occurrence_set_id' in its params.") + return None # Scoped to the project, since params are set by any member who can create a job. occurrence_set = OccurrenceSet.objects.for_project(job.project).filter(pk=occurrence_set_id).first() if not occurrence_set: diff --git a/ami/ml/test_training.py b/ami/ml/test_training.py index bf4c8fad9..25b2c780e 100644 --- a/ami/ml/test_training.py +++ b/ami/ml/test_training.py @@ -260,7 +260,7 @@ def test_someone_who_cannot_retrain_cannot_read_the_training_set(self): self.assertEqual(response.status_code, 403) -class TestTheJobRequiresASet(TrainingSetFixture): +class TestTheJobsOccurrenceSet(TrainingSetFixture): def _job(self, **params): from ami.jobs.models import Job, TrainClassifierJob @@ -271,13 +271,10 @@ def _job(self, **params): params={"algorithm_key": self.algorithm.key, **params}, ) - def test_a_job_without_a_set_is_refused(self): + def test_a_job_without_a_set_learns_from_everything_verified(self): from ami.jobs.models import TrainClassifierJob - with self.assertRaises(ValueError) as caught: - TrainClassifierJob.target_occurrence_set(self._job()) - - self.assertIn("occurrence_set_id", str(caught.exception)) + self.assertIsNone(TrainClassifierJob.target_occurrence_set(self._job())) def test_a_set_from_another_project_is_refused(self): from ami.jobs.models import TrainClassifierJob @@ -297,3 +294,4 @@ def test_a_set_in_this_project_is_used(self): mine = self.make_set("Mine", [self.make_occurrence(self.taxa[0])]) self.assertEqual(TrainClassifierJob.target_occurrence_set(self._job(occurrence_set_id=mine.pk)), mine) + From 476cf88db58b3613ad7a8851613d2ab81882ed50 Mon Sep 17 00:00:00 2001 From: mohamedelabbas1996 Date: Thu, 8 Oct 2026 17:08:05 -0400 Subject: [PATCH 4/8] fix(jobs): retrain the newest version of the algorithm named A key is shared by every version of an algorithm, while embeddings are stored against one version row. Taking whichever row the database returned first could pick an older version and train on an empty set of vectors, or on vectors from a different fitting, with nothing in the job to say which had happened. --- ami/jobs/models.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/ami/jobs/models.py b/ami/jobs/models.py index 49505f81d..f0789dc8b 100644 --- a/ami/jobs/models.py +++ b/ami/jobs/models.py @@ -1037,7 +1037,9 @@ def run(cls, job: "Job"): if not algorithm_key: raise ValueError("A train_classifier job needs an 'algorithm_key' in its params.") - algorithm = Algorithm.objects.filter(key=algorithm_key).first() + # A key is shared by every version of an algorithm, and embeddings are stored per + # version row, so the newest is chosen rather than whichever row came back first. + algorithm = Algorithm.objects.filter(key=algorithm_key).order_by("-version", "-pk").first() if not algorithm: raise ValueError(f"No algorithm with key '{algorithm_key}'.") if not algorithm.trainable: From 73a49db32215791524c71e4f445299967180778c Mon Sep 17 00:00:00 2001 From: mohamedelabbas1996 Date: Thu, 8 Oct 2026 19:25:34 -0400 Subject: [PATCH 5/8] feat(jobs): follow a training run's epochs while it fits Training is the long stage of the job and the service says nothing while it runs, so the bar sat at nought for as long as the fitting took and a person could not tell a slow run from a dead one. The service is now told where to report, and posts the epoch it has reached. The callback is authorised by the same signed token as the result and the head upload, since a service has no Antenna account. The epoch and the total land as stage parameters, and the stage progress is the ratio of the two. The total is written at dispatch from the run's own settings, so the stage reads "0 of 5000" while the service is still starting rather than an empty pair. A ping never finishes the stage: the result does, because the service still has to score the head and upload it after the last epoch. Pings that arrive late or out of order are dropped rather than dragging a run backwards. Reporting is optional. A service that does not post anything behaves exactly as before, and a ping that fails to send must not fail the run. --- ami/jobs/models.py | 54 ++++++++++++++++++++++++ ami/jobs/views.py | 41 ++++++++++++++++++ ami/ml/test_training.py | 84 +++++++++++++++++++++++++++++++++++++ ami/ml/training/__init__.py | 2 + ami/ml/training/service.py | 9 ++++ 5 files changed, 190 insertions(+) diff --git a/ami/jobs/models.py b/ami/jobs/models.py index f0789dc8b..5fb3bfdd8 100644 --- a/ami/jobs/models.py +++ b/ami/jobs/models.py @@ -1026,6 +1026,10 @@ class TrainClassifierJob(JobType): STAGE_DISPATCH = "dispatch" STAGE_TRAIN = "train" + # Reported by the service while it fits, and read back to move the training stage. + PARAM_EPOCH = "Epoch" + PARAM_TOTAL_EPOCHS = "Total epochs" + @classmethod def run(cls, job: "Job"): from ami.ml.models import Algorithm @@ -1204,6 +1208,12 @@ def dispatch(cls, job: "Job", service, algorithm, dataset: dict) -> None: """Hand the service the dataset URL and leave the job running until it reports back.""" from ami.ml.training import send_training_request + # Written before the request so the stage reads "0 of 300" while the service is + # still starting, rather than an empty pair of parameters. + config = cls.config_for(job, algorithm) + job.progress.add_stage_param(cls.STAGE_TRAIN, cls.PARAM_EPOCH, 0) + job.progress.add_stage_param(cls.STAGE_TRAIN, cls.PARAM_TOTAL_EPOCHS, config.epochs) + response = send_training_request(job=job, service=service, algorithm=algorithm, dataset=dataset) # The service posts its callback before returning from /train, so a fast run is @@ -1240,6 +1250,50 @@ def dispatch(cls, job: "Job", service, algorithm, dataset: dict) -> None: }, ) + @classmethod + def record_progress(cls, job: "Job", epoch: int, total_epochs: int | None = None) -> None: + """ + Move the training stage to the epoch a service says it has reached. + + Training is the long stage and the service is silent while it runs, so without this + the stage sits at nought for as long as the run takes. The epoch count is the only + progress the service has to report. + """ + # As in record_result: a callback can arrive on a job whose stages were never set + # up, for instance after a restart. + if not any(stage.key == cls.STAGE_TRAIN for stage in job.progress.stages): + job.progress.add_stage("Training", cls.STAGE_TRAIN) + + try: + reached = job.progress.get_stage_param(cls.STAGE_TRAIN, job.progress.make_key(cls.PARAM_EPOCH)).value + except ValueError: + reached = 0 + # Pings can arrive late or out of order, and a stale one must not drag a run + # backwards past the epoch it actually reached. + if epoch <= (reached or 0): + return + + total = total_epochs + if not total: + try: + total = job.progress.get_stage_param( + cls.STAGE_TRAIN, job.progress.make_key(cls.PARAM_TOTAL_EPOCHS) + ).value + except ValueError: + total = None + + job.progress.add_or_update_stage_param(cls.STAGE_TRAIN, cls.PARAM_EPOCH, epoch) + if total: + job.progress.add_or_update_stage_param(cls.STAGE_TRAIN, cls.PARAM_TOTAL_EPOCHS, total) + job.progress.update_stage( + cls.STAGE_TRAIN, + status=JobState.STARTED, + # Never a full stage: the result finishes it, not the last epoch, since the + # service still has to score the head and upload it. + progress=min(epoch / total, 0.99) if total else 0, + ) + job.save(update_fields=["progress", "updated_at"]) + @classmethod def record_result(cls, job: "Job", payload: dict) -> None: """Store what the service reported, register the new version, and finish the job.""" diff --git a/ami/jobs/views.py b/ami/jobs/views.py index 62cc64dc9..449271dd2 100644 --- a/ami/jobs/views.py +++ b/ami/jobs/views.py @@ -653,6 +653,47 @@ def training_result(self, request, pk=None): logger.info("Recorded a training result for job %s", job.pk) return Response({"status": "recorded"}) + @extend_schema(exclude=True) + @action( + detail=True, + methods=["post"], + url_path="training-progress", + name="training-progress", + # Same token as the result callback: a processing service has no Antenna account. + permission_classes=[AllowAny], + authentication_classes=[], + ) + def training_progress(self, request, pk=None): + """ + Receive how far through its epochs a retraining run is. + + Training is the long stage of the job and the service says nothing while it fits, so + this is what moves the progress bar in between. + """ + from ami.jobs.models import TrainClassifierJob + + job = self._job_for_callback(pk, request) + + if job.job_type_key != TrainClassifierJob.key: + raise ValidationError(f"Job #{job.pk} is not a training job.") + + if job.status in JobState.final_states(): + # A ping overtaken by the result. The job is done; its stages say so. + return Response({"status": "finished"}) + + epoch = SingleParamSerializer[int].clean( + "epoch", + serializers.IntegerField(required=True, min_value=0), + request.data, + ) + total_epochs = SingleParamSerializer[int].clean( + "total_epochs", + serializers.IntegerField(required=False, allow_null=True, default=None, min_value=1), + request.data, + ) + TrainClassifierJob.record_progress(job=job, epoch=epoch, total_epochs=total_epochs) + return Response({"status": "recorded"}) + @extend_schema(exclude=True) @action( detail=True, diff --git a/ami/ml/test_training.py b/ami/ml/test_training.py index 25b2c780e..659b2ac96 100644 --- a/ami/ml/test_training.py +++ b/ami/ml/test_training.py @@ -260,6 +260,7 @@ def test_someone_who_cannot_retrain_cannot_read_the_training_set(self): self.assertEqual(response.status_code, 403) +<<<<<<< HEAD class TestTheJobsOccurrenceSet(TrainingSetFixture): def _job(self, **params): from ami.jobs.models import Job, TrainClassifierJob @@ -295,3 +296,86 @@ def test_a_set_in_this_project_is_used(self): self.assertEqual(TrainClassifierJob.target_occurrence_set(self._job(occurrence_set_id=mine.pk)), mine) +======= +class TestTheTrainingStageFollowsTheEpochs(TrainingSetFixture): + """ + Training is the long stage and the service is silent while it fits, so the epochs it + reports are the only thing that can move the stage while a run is going. + """ + + def setUp(self) -> None: + super().setUp() + from ami.jobs.models import Job, TrainClassifierJob + + self.job = Job.objects.create( + project=self.project, + name="Retrain", + job_type_key=TrainClassifierJob.key, + params={"algorithm_key": self.algorithm.key}, + ) + self.job.progress.add_stage("Training", TrainClassifierJob.STAGE_TRAIN) + self.url = reverse("api:job-training-progress", args=[self.job.pk]) + + def _stage(self): + from ami.jobs.models import Job, TrainClassifierJob + + job = Job.objects.get(pk=self.job.pk) + return job.progress.get_stage(TrainClassifierJob.STAGE_TRAIN) + + def _param(self, name: str): + from ami.jobs.models import Job, TrainClassifierJob + + job = Job.objects.get(pk=self.job.pk) + return job.progress.get_stage_param(TrainClassifierJob.STAGE_TRAIN, job.progress.make_key(name)).value + + def _post(self, **payload): + from ami.ml.training import make_callback_token + + return APIClient().post( + self.url, + payload, + format="json", + HTTP_AUTHORIZATION=f"Token {make_callback_token(self.job)}", + ) + + def test_a_ping_moves_the_stage(self): + from ami.jobs.models import TrainClassifierJob + + response = self._post(epoch=150, total_epochs=300) + + self.assertEqual(response.status_code, 200, response.data) + self.assertEqual(self._param(TrainClassifierJob.PARAM_EPOCH), 150) + self.assertEqual(self._param(TrainClassifierJob.PARAM_TOTAL_EPOCHS), 300) + self.assertEqual(self._stage().progress, 0.5) + + def test_the_stage_is_never_finished_by_an_epoch(self): + """The result finishes it: the service still has to score the head and upload it.""" + self._post(epoch=300, total_epochs=300) + + self.assertEqual(self._stage().progress, 0.99) + + def test_a_ping_that_arrives_late_is_ignored(self): + self._post(epoch=200, total_epochs=300) + self._post(epoch=100, total_epochs=300) + + from ami.jobs.models import TrainClassifierJob + + self.assertEqual(self._param(TrainClassifierJob.PARAM_EPOCH), 200) + + def test_a_ping_without_the_token_is_refused(self): + response = APIClient().post(self.url, {"epoch": 10}, format="json") + + self.assertEqual(response.status_code, 403) + + def test_a_ping_after_the_job_finished_changes_nothing(self): + from ami.jobs.models import JobState, TrainClassifierJob + + self.job.status = JobState.SUCCESS + self.job.save() + + response = self._post(epoch=10, total_epochs=300) + + self.assertEqual(response.status_code, 200) + with self.assertRaises(ValueError): + self._param(TrainClassifierJob.PARAM_EPOCH) +>>>>>>> ef15edc4 (feat(jobs): follow a training run's epochs while it fits) diff --git a/ami/ml/training/__init__.py b/ami/ml/training/__init__.py index 268c6c6d5..91bab508e 100644 --- a/ami/ml/training/__init__.py +++ b/ami/ml/training/__init__.py @@ -32,6 +32,7 @@ head_path, head_upload_url_for, make_callback_token, + progress_url_for, send_training_request, store_head, verify_callback_token, @@ -56,6 +57,7 @@ "head_upload_url_for", "label_counts", "make_callback_token", + "progress_url_for", "row_as_dict", "send_training_request", "species_with_enough_examples", diff --git a/ami/ml/training/service.py b/ami/ml/training/service.py index b9a7c4920..da9866a9c 100644 --- a/ami/ml/training/service.py +++ b/ami/ml/training/service.py @@ -89,6 +89,12 @@ def head_upload_url_for(job: "Job") -> str: return urljoin(_callback_base(job).rstrip("/") + "/", path.lstrip("/")) +def progress_url_for(job: "Job") -> str: + """Where the service should report how far through the epochs it is.""" + path = reverse("api:job-training-progress", args=[job.pk]) + return urljoin(_callback_base(job).rstrip("/") + "/", path.lstrip("/")) + + def absolute_media_url(url: str, base_url: str | None = None) -> str: """ Turn a stored file's URL into one a processing service can fetch. @@ -143,6 +149,9 @@ def send_training_request( # The head itself comes back here. A service that does not support the upload simply # ignores this, and the version is registered without a stored copy as before. payload["head_upload_url"] = head_upload_url_for(job) + # Optional for the service: without it the training stage simply stays at nought until + # the result lands, as it did before. + payload["progress_url"] = progress_url_for(job) job.logger.info(f"Sending training request to {endpoint} for {algorithm.key}") session = create_session() From b01bc18440b23ddade58bc9a2cb1fa20be05b199 Mon Sep 17 00:00:00 2001 From: mohamedelabbas1996 Date: Fri, 9 Oct 2026 00:32:37 -0400 Subject: [PATCH 6/8] feat(ml): write down the contract between a training job and a service The training request was a dict literal and the result was read with .get() at a dozen call sites, so neither side of the exchange was written down anywhere. Every other service contract in this module is a schema; this one was the exception, and a service renaming a field would have produced an empty metrics dict and no error. TrainingRequest and TrainingResult mirror the service's own TrainRequest and TrainResponse, the way PipelineRequest and PipelineResultsResponse already do for processing. TrainingResult keeps unknown fields: a newer service may report more than this one knows about, and the result is the only record of what a run did. incumbent_metrics is nullable because a service reports null when there was no current head to score, which is exactly what a first retrain does. --- ami/ml/schemas.py | 58 +++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 58 insertions(+) diff --git a/ami/ml/schemas.py b/ami/ml/schemas.py index d1d00a575..88ff797af 100644 --- a/ami/ml/schemas.py +++ b/ami/ml/schemas.py @@ -473,6 +473,64 @@ class AsyncPipelineRegistrationRequest(pydantic.BaseModel): pipelines: list[PipelineConfigResponse] = [] +class TrainingRequest(pydantic.BaseModel): + """ + What Antenna asks a processing service to do, mirroring the service's own TrainRequest. + + Declared rather than assembled as a dict so the two sides of the contract are written + down in one place, the way the pipeline request and response already are. + """ + + dataset_url: str = pydantic.Field(description="Where the service downloads the training set from.") + algorithm_key: str = pydantic.Field(description="The head to retrain. Its current weights are the baseline.") + job_id: int + name: str + + # The settings the training set was built under, plus the ones the service fits with. + min_per_species: int + min_improvement: float + head_type: str + epochs: int + learning_rate: float + weight_decay: float + + # Where the service reports back. The result callback is the only way a run that + # outlasts its request can finish, so it is always sent; the other two are optional for + # the service and ignored by one that does not support them. + callback_url: str + callback_token: str + head_upload_url: str | None = None + progress_url: str | None = None + + +class TrainingResult(pydantic.BaseModel): + """ + What a processing service reports when a run finishes, mirroring its TrainResponse. + + Extra fields are kept: a newer service may report more than this one knows about, and + dropping it silently would lose the only record of what a run did. + """ + + promote: bool = False + reason: str = "" + warnings: list[str] = pydantic.Field(default_factory=list) + labels: list[str] = pydantic.Field(default_factory=list, description="The new head's classes, in order.") + rows: dict = pydantic.Field(default_factory=dict, description="How many rows were used, kept and dropped.") + counts: dict = pydantic.Field(default_factory=dict) + dropped_species: list[str] = pydantic.Field(default_factory=list) + classes_restored_from_current_head: int = 0 + candidate_metrics: dict = pydantic.Field(default_factory=dict, description="The new head, on the held-out rows.") + # Null when there was no current head to score, which a service reports rather than + # leaving out, so the field has to accept it. + incumbent_metrics: dict | None = pydantic.Field(default=None, description="The current head, on the same rows.") + saved: dict[str, str] | None = None + head_url: str | None = pydantic.Field(default=None, description="Where the uploaded head landed, if it was sent.") + trained_at: str | None = None + + class Config: + extra = "allow" + + class AlgorithmTrainingInfo(pydantic.BaseModel): """ What actually happened when this version was trained. Written by the service, read-only. From 7faf0c13b8f291c0177122293f62f7be949ca3b4 Mon Sep 17 00:00:00 2001 From: mohamedelabbas1996 Date: Fri, 9 Oct 2026 00:37:19 -0400 Subject: [PATCH 7/8] fix(jobs): take a training result one way only, so there is one shape to read A result could arrive two ways: in the body of the /train response, or at the job's callback. The two carried the same outcome in different shapes, and both of these bugs lived in that gap. On the callback path the dataset metadata the service echoes was read as payload["dataset"]["metadata"], which is not where it is, so it was never parsed and a registered version never recorded which occurrence set it learned from. The provenance this is all for was quietly dropped on every real run. On the inline path the same code would have raised AttributeError, because there "dataset" is the schema object rather than a dict. It never fired only because a service posts its callback first and the duplicate guard returned early. /train is now an acknowledgement: the result arrives at the callback, is validated by a serializer at the view, and record_result takes the parsed result rather than a raw payload. One shape, one path, parsed once at the boundary. The result is validated because the new algorithm version is built from it. The echoed dataset metadata is read leniently: a service echoing an older shape should still have its result recorded, and the cost is a version that cannot say which set it came from, not a head that is lost. Also fixes a flaky test: with a handful of rows, a hash split does not give a particular ratio a particular count, so the ratio is now compared across two values rather than asserted exactly. --- ami/jobs/models.py | 118 ++++++++++++++--------------- ami/jobs/serializers.py | 37 ++++++++- ami/jobs/views.py | 54 +++++++------ ami/ml/test_training.py | 151 +++++++++++++++++++++++++++++++++++-- ami/ml/training/service.py | 64 +++++++--------- 5 files changed, 300 insertions(+), 124 deletions(-) diff --git a/ami/jobs/models.py b/ami/jobs/models.py index 5fb3bfdd8..10ebcca2a 100644 --- a/ami/jobs/models.py +++ b/ami/jobs/models.py @@ -24,6 +24,10 @@ logger = logging.getLogger(__name__) +if typing.TYPE_CHECKING: + from ami.ml.models import Algorithm + from ami.ml.schemas import TrainingResult + class JobDispatchMode(models.TextChoices): """ @@ -1214,11 +1218,11 @@ def dispatch(cls, job: "Job", service, algorithm, dataset: dict) -> None: job.progress.add_stage_param(cls.STAGE_TRAIN, cls.PARAM_EPOCH, 0) job.progress.add_stage_param(cls.STAGE_TRAIN, cls.PARAM_TOTAL_EPOCHS, config.epochs) - response = send_training_request(job=job, service=service, algorithm=algorithm, dataset=dataset) + send_training_request(job=job, service=service, algorithm=algorithm, dataset=dataset) - # The service posts its callback before returning from /train, so a fast run is - # already finished by the time this line is reached. Writing this copy's progress - # would put the training stage back to "started" on a job that is done. + # A fast run reports its result before the request it came from returns, so the + # job can already be finished here. Writing this copy's progress would put the + # training stage back to "started" on a job that is done. job.refresh_from_db() if job.status in JobState.final_states(): job.logger.info(f"{service.name} already reported its result.") @@ -1226,29 +1230,10 @@ def dispatch(cls, job: "Job", service, algorithm, dataset: dict) -> None: job.progress.update_stage(cls.STAGE_DISPATCH, status=JobState.SUCCESS, progress=1) job.progress.update_stage(cls.STAGE_TRAIN, status=JobState.STARTED, progress=0) - job.logger.info(f"Training request accepted by {service.name}. Waiting for it to report back.") job.save(update_fields=["progress", "updated_at"]) - - if response is None: - # The service accepted the work and will post to the job's training-result - # endpoint when it finishes. Leaving the job STARTED is the point: a real - # training set outlasts the request that started it. - job.logger.info("Waiting for the service to report back.") - return - - if response is not None: - # Answered inline, which small datasets do. Record it now rather than waiting - # for a callback that has already been overtaken. - cls.record_result( - job=job, - payload={ - "result": response, - "dataset": dataset["metadata"], - # Kept alongside the metadata so the registered version can point at the - # exact file it was fitted on, not just describe it. - "dataset_url": dataset["url"], - }, - ) + # Leaving the job STARTED is the point: a real training set outlasts the request + # that started it, and the result arrives at the job's callback. + job.logger.info(f"Training request accepted by {service.name}. Waiting for it to report back.") @classmethod def record_progress(cls, job: "Job", epoch: int, total_epochs: int | None = None) -> None: @@ -1295,19 +1280,28 @@ def record_progress(cls, job: "Job", epoch: int, total_epochs: int | None = None job.save(update_fields=["progress", "updated_at"]) @classmethod - def record_result(cls, job: "Job", payload: dict) -> None: - """Store what the service reported, register the new version, and finish the job.""" - # A service posts its callback before returning from /train, so a fast run reports - # twice: once through the callback and once inline. Without this guard each retrain - # registered two algorithm versions. Re-read first, because the inline caller holds - # a copy from before the callback landed. + def record_result( + cls, + job: "Job", + result: "TrainingResult", + dataset: dict | None = None, + dataset_url: str | None = None, + ) -> "Algorithm | None": + """ + Store what the service reported, register the new version, and finish the job. + + The result arrives one way only: through the callback, parsed at the view. A + service that also answers inline is acknowledged and its body ignored, so there is + a single shape to read here. + """ job.refresh_from_db(fields=["status", "result"]) if job.status in JobState.final_states(): + # A retry of the callback, or a second service reporting on the same job. The + # first answer stands; recording this one would register a second version. job.logger.info("A training result is already recorded for this job; ignoring a duplicate.") - return + return None - result = payload.get("result") or {} - job.result = payload + job.result = {"result": result.dict(), "dataset": dataset, "dataset_url": dataset_url} # A result can arrive through the callback on a job whose stages were never set up, # for instance after a restart. Make sure the stage exists before reporting into it. @@ -1318,17 +1312,16 @@ def record_result(cls, job: "Job", payload: dict) -> None: if any(stage.key == cls.STAGE_DISPATCH for stage in job.progress.stages): job.progress.update_stage(cls.STAGE_DISPATCH, status=JobState.SUCCESS, progress=1) - for warning in result.get("warnings", []): + for warning in result.warnings: job.logger.warning(warning) - candidate = result.get("candidate_metrics") or {} - incumbent = result.get("incumbent_metrics") or {} - job.progress.add_stage_param(cls.STAGE_TRAIN, "New head top-1", candidate.get("top1")) + job.progress.add_stage_param(cls.STAGE_TRAIN, "New head top-1", result.candidate_metrics.get("top1")) + incumbent = result.incumbent_metrics or {} job.progress.add_stage_param(cls.STAGE_TRAIN, "Current head top-1", incumbent.get("top1")) - job.progress.add_stage_param(cls.STAGE_TRAIN, "Better", result.get("promote")) - job.logger.info(result.get("reason", "Training finished.")) + job.progress.add_stage_param(cls.STAGE_TRAIN, "Better", result.promote) + job.logger.info(result.reason or "Training finished.") - new_version = cls.register_new_version(job=job, payload=payload) + new_version = cls.register_new_version(job=job, result=result, dataset=dataset, dataset_url=dataset_url) if new_version: job.progress.add_stage_param(cls.STAGE_TRAIN, "New version", new_version.key) job.logger.info(f"Registered {new_version} as version {new_version.version}") @@ -1337,9 +1330,16 @@ def record_result(cls, job: "Job", payload: dict) -> None: job.finished_at = datetime.datetime.now() job.update_status(JobState.SUCCESS, save=True) job.save() + return new_version @classmethod - def register_new_version(cls, job: "Job", payload: dict): + def register_new_version( + cls, + job: "Job", + result: "TrainingResult", + dataset: dict | None = None, + dataset_url: str | None = None, + ): """ Record the trained head as a new version of the algorithm it was trained from. @@ -1352,29 +1352,29 @@ def register_new_version(cls, job: "Job", payload: dict): from ami.ml.models import Algorithm, AlgorithmCategoryMap from ami.ml.schemas import AlgorithmTrainingInfo, TrainingDatasetMetadata - result = payload.get("result") or {} - dataset = payload.get("dataset") or {} - # The service echoes the dataset's metadata back as JSON, so it arrives as a dict - # rather than the model that wrote it. Parsed here so a malformed echo is caught at - # the boundary instead of further in. + # The service echoes the training set's own metadata, which is where the run's + # provenance comes from. Read leniently: a service echoing an older shape should + # still have its result recorded, and the only cost is a version that cannot say + # which set it learned from. metadata = None - if dataset.get("metadata"): + if dataset: try: - metadata = TrainingDatasetMetadata.parse_obj(dataset["metadata"]) + metadata = TrainingDatasetMetadata.parse_obj(dataset) except pydantic.ValidationError as e: job.logger.warning(f"The service returned training set metadata Antenna could not read: {e}") + parent_key = (job.params or {}).get("algorithm_key") parent = Algorithm.objects.filter(key=parent_key).first() if not parent: job.logger.warning(f"Cannot register a new version: no algorithm with key '{parent_key}'.") return None - labels = result.get("labels") or [] + labels = result.labels if not labels: job.logger.warning("The service returned no class list, so no new version was registered.") return None - trained_at = result.get("trained_at") or "" + trained_at = result.trained_at or "" stamp = str(trained_at).replace(":", "").replace("-", "").replace(".", "")[:15] or f"job{job.pk}" version = ( Algorithm.objects.filter(name=parent.name).order_by("-version").values_list("version", flat=True).first() @@ -1386,7 +1386,7 @@ def register_new_version(cls, job: "Job", payload: dict): data=[{"label": label, "index": i} for i, label in enumerate(labels)], version=stamp, description=( - f"Retrained from {dataset.get('rows', len(labels))} verified crops in " + f"Retrained from {metadata.rows if metadata else len(labels)} verified crops in " f"project '{job.project.name}' by job #{job.pk}." ), ) @@ -1396,7 +1396,7 @@ def register_new_version(cls, job: "Job", payload: dict): key=f"{parent.key}-v{version}-{stamp}", # Where the weights are kept. Empty when the service did not upload them, in # which case the head exists only on that service's own disk. - uri=result.get("head_url") or "", + uri=result.head_url or "", version=version, version_name=str(trained_at) or stamp, task_type=parent.task_type, @@ -1406,16 +1406,16 @@ def register_new_version(cls, job: "Job", payload: dict): category_map=category_map, training_info=AlgorithmTrainingInfo( trained_at=trained_at or None, - dataset_url=payload.get("dataset_url") or dataset.get("url"), + dataset_url=dataset_url or (metadata.url if metadata else None), occurrence_set_id=metadata.occurrence_set.id if metadata and metadata.occurrence_set else None, occurrence_set_name=metadata.occurrence_set.name if metadata and metadata.occurrence_set else None, - dataset_rows=(result.get("rows") or {}).get("kept"), + dataset_rows=result.rows.get("kept"), dataset_classes=len(labels), - metrics=result.get("candidate_metrics") or {}, - previous_metrics=result.get("incumbent_metrics") or {}, + metrics=result.candidate_metrics, + previous_metrics=result.incumbent_metrics or {}, parent_algorithm_key=parent.key, job_id=job.pk, - warnings=result.get("warnings") or [], + warnings=result.warnings, ), ) diff --git a/ami/jobs/serializers.py b/ami/jobs/serializers.py index 7fe1e2581..ac521f39e 100644 --- a/ami/jobs/serializers.py +++ b/ami/jobs/serializers.py @@ -11,7 +11,7 @@ ) from ami.main.models import Deployment, Project, SourceImage, SourceImageCollection from ami.ml.models import Pipeline -from ami.ml.schemas import PipelineProcessingTask, PipelineTaskResult, ProcessingServiceClientInfo +from ami.ml.schemas import PipelineProcessingTask, PipelineTaskResult, ProcessingServiceClientInfo, TrainingResult from ami.ml.serializers import PipelineNestedSerializer from .models import ( @@ -292,6 +292,41 @@ class MLJobResultsRequestSerializer(serializers.Serializer): client_info = SchemaField(schema=ProcessingServiceClientInfo, required=False, default=None) +class TrainingResultRequestSerializer(serializers.Serializer): + """POST /jobs/{id}/training-result/ — the body a processing service posts when a run finishes. + + The counterpart of MLJobResultsRequestSerializer for a retraining job: one result for + one job, rather than a list of per-item results. + + ``result`` is validated, since it is what the new algorithm version is built from. + ``dataset`` is the metadata the service echoes back from the training set file and is + taken as it comes: a service that echoes an older shape should still have its result + recorded, and the cost is a version that cannot say which occurrence set it learned + from. It is parsed, leniently, where it is read. + """ + + result = SchemaField(schema=TrainingResult) + dataset = serializers.JSONField(required=False, allow_null=True, default=None) + dataset_url = serializers.CharField(required=False, allow_null=True, default=None) + job_id = serializers.IntegerField(required=False, allow_null=True, default=None) + algorithm_key = serializers.CharField(required=False, allow_null=True, default=None) + + +class TrainingResultResponseSerializer(serializers.Serializer): + """POST /jobs/{id}/training-result/ — acknowledgment returned to the processing service.""" + + status = serializers.CharField() + job_id = serializers.IntegerField() + algorithm = serializers.CharField(allow_null=True, help_text="Key of the version registered, if one was.") + + +class TrainingProgressRequestSerializer(serializers.Serializer): + """POST /jobs/{id}/training-progress/ — how far through its epochs a run has got.""" + + epoch = serializers.IntegerField(min_value=0) + total_epochs = serializers.IntegerField(required=False, allow_null=True, default=None, min_value=1) + + class MLJobResultsResponseSerializer(serializers.Serializer): """POST /jobs/{id}/result/ — acknowledgment returned to the processing service. diff --git a/ami/jobs/views.py b/ami/jobs/views.py index 449271dd2..175ea70b4 100644 --- a/ami/jobs/views.py +++ b/ami/jobs/views.py @@ -31,6 +31,9 @@ MLJobResultsResponseSerializer, MLJobTasksRequestSerializer, MLJobTasksResponseSerializer, + TrainingProgressRequestSerializer, + TrainingResultRequestSerializer, + TrainingResultResponseSerializer, ) from ami.jobs.tasks import ( HEARTBEAT_THROTTLE_SECONDS, @@ -614,11 +617,9 @@ def _job_for_callback(self, pk, request) -> Job: return job @extend_schema( - request=MLJobResultsRequestSerializer, - responses={200: MLJobResultsResponseSerializer}, - parameters=[project_id_doc_param], + request=TrainingResultRequestSerializer, + responses={200: TrainingResultResponseSerializer}, ) - @extend_schema(exclude=True) @action( detail=True, methods=["post"], @@ -633,8 +634,9 @@ def training_result(self, request, pk=None): """ Receive the outcome of a retraining run from a processing service. - Training can outlast the request that started it, so the service reports back here - instead of holding the connection open. + Training outlasts the request that started it, so this is the one way a result + comes back. A service that also answers its caller inline is acknowledged and that + body ignored, so there is a single shape to read. """ from ami.jobs.models import TrainClassifierJob @@ -644,14 +646,27 @@ def training_result(self, request, pk=None): raise ValidationError(f"Job #{job.pk} is not a training job.") if job.status in JobState.final_states(): - # The service answered inline and the result is already recorded, or a retry - # arrived late. Either way the first answer stands. + # A retry of a callback that already landed. The first answer stands. logger.info("Ignoring a training result for job %s, which already finished", job.pk) - return Response({"status": "already recorded"}) + return Response({"status": "already recorded", "job_id": job.pk, "algorithm": None}) + + serializer = TrainingResultRequestSerializer(data=request.data) + serializer.is_valid(raise_exception=True) - TrainClassifierJob.record_result(job=job, payload=request.data) + new_version = TrainClassifierJob.record_result( + job=job, + result=serializer.validated_data["result"], + dataset=serializer.validated_data["dataset"], + dataset_url=serializer.validated_data["dataset_url"], + ) logger.info("Recorded a training result for job %s", job.pk) - return Response({"status": "recorded"}) + return Response( + { + "status": "recorded", + "job_id": job.pk, + "algorithm": new_version.key if new_version else None, + } + ) @extend_schema(exclude=True) @action( @@ -681,17 +696,14 @@ def training_progress(self, request, pk=None): # A ping overtaken by the result. The job is done; its stages say so. return Response({"status": "finished"}) - epoch = SingleParamSerializer[int].clean( - "epoch", - serializers.IntegerField(required=True, min_value=0), - request.data, - ) - total_epochs = SingleParamSerializer[int].clean( - "total_epochs", - serializers.IntegerField(required=False, allow_null=True, default=None, min_value=1), - request.data, + serializer = TrainingProgressRequestSerializer(data=request.data) + serializer.is_valid(raise_exception=True) + + TrainClassifierJob.record_progress( + job=job, + epoch=serializer.validated_data["epoch"], + total_epochs=serializer.validated_data["total_epochs"], ) - TrainClassifierJob.record_progress(job=job, epoch=epoch, total_epochs=total_epochs) return Response({"status": "recorded"}) @extend_schema(exclude=True) diff --git a/ami/ml/test_training.py b/ami/ml/test_training.py index 659b2ac96..b75eede40 100644 --- a/ami/ml/test_training.py +++ b/ami/ml/test_training.py @@ -217,15 +217,22 @@ def test_a_set_from_another_project_is_not_counted(self): self.assertEqual(response.status_code, 404) def test_the_split_ratio_moves_the_counts(self): + """ + Compared across two ratios rather than asserted exactly. + + Which side a given occurrence falls on is a hash of its id, so with a handful of + rows a particular ratio does not produce a particular count. + """ for taxon in self.taxa: self.make_occurrence(taxon) self.make_occurrence(taxon) - everything_held_out = self._summary(test_fraction=0.99) + mostly_train = self._summary(test_fraction=0.01) + mostly_test = self._summary(test_fraction=0.99) - self.assertEqual(everything_held_out["test"], everything_held_out["rows"]) - self.assertEqual(everything_held_out["train"], 0) - self.assertEqual(everything_held_out["settings"]["test_fraction"], 0.99) + self.assertEqual(mostly_test["rows"], mostly_train["rows"]) + self.assertGreater(mostly_test["test"], mostly_train["test"]) + self.assertEqual(mostly_test["settings"]["test_fraction"], 0.99) def test_a_ratio_no_run_could_use_is_refused(self): response = self.client.get( @@ -260,7 +267,6 @@ def test_someone_who_cannot_retrain_cannot_read_the_training_set(self): self.assertEqual(response.status_code, 403) -<<<<<<< HEAD class TestTheJobsOccurrenceSet(TrainingSetFixture): def _job(self, **params): from ami.jobs.models import Job, TrainClassifierJob @@ -296,7 +302,7 @@ def test_a_set_in_this_project_is_used(self): self.assertEqual(TrainClassifierJob.target_occurrence_set(self._job(occurrence_set_id=mine.pk)), mine) -======= + class TestTheTrainingStageFollowsTheEpochs(TrainingSetFixture): """ Training is the long stage and the service is silent while it fits, so the epochs it @@ -378,4 +384,135 @@ def test_a_ping_after_the_job_finished_changes_nothing(self): self.assertEqual(response.status_code, 200) with self.assertRaises(ValueError): self._param(TrainClassifierJob.PARAM_EPOCH) ->>>>>>> ef15edc4 (feat(jobs): follow a training run's epochs while it fits) + + +class TestTheResultComingBackFromAService(TrainingSetFixture): + """ + The callback is the one way a result arrives, so its shape is the contract. + + A service posts the result together with the training set's own metadata, echoed back + from the file Antenna wrote. That echo is where a registered version learns which + occurrence set it came from. + """ + + def setUp(self) -> None: + super().setUp() + from ami.jobs.models import Job, TrainClassifierJob + + self.occurrence_set = self.make_set("Mine", [self.make_occurrence(self.taxa[0])]) + self.job = Job.objects.create( + project=self.project, + name="Retrain", + job_type_key=TrainClassifierJob.key, + params={"algorithm_key": self.algorithm.key, "occurrence_set_id": self.occurrence_set.pk}, + ) + self.url = reverse("api:job-training-result", args=[self.job.pk]) + + def _dataset_echo(self, **overrides) -> dict: + """The metadata a service reads out of the training set and sends back.""" + from ami.ml.schemas import ( + NamedReference, + TrainedAlgorithmReference, + TrainingDatasetMetadata, + TrainingDatasetSettings, + ) + + metadata = TrainingDatasetMetadata( + url="/media/training/set.npz", + project=NamedReference(id=self.project.pk, name=self.project.name), + algorithm=TrainedAlgorithmReference(key=self.algorithm.key, name=self.algorithm.name, version=1), + dimensions=VECTOR_LENGTH, + dtype="float32", + classes=[t.name for t in self.taxa], + rows=6, + train=5, + test=1, + occurrence_set=NamedReference(id=self.occurrence_set.pk, name=self.occurrence_set.name), + settings=TrainingDatasetSettings(min_per_species=2, split_salt="antenna-head-v1", test_fraction=0.2), + ) + return {**metadata.dict(), **overrides} + + def _post(self, payload: dict): + from ami.ml.training import make_callback_token + + return APIClient().post( + self.url, + payload, + format="json", + HTTP_AUTHORIZATION=f"Token {make_callback_token(self.job)}", + ) + + def _result(self, **overrides) -> dict: + return { + "promote": False, + "reason": "The new head did not beat the current one.", + "warnings": ["Only 1 held-out row(s)."], + "labels": [t.name for t in self.taxa], + "rows": {"kept": 6, "dropped": 0}, + "candidate_metrics": {"top1": 0.8}, + "incumbent_metrics": {"top1": 0.9}, + "trained_at": "2026-10-08T23:00:00", + "head_url": "/media/algorithms/head.npz", + **overrides, + } + + def test_a_result_registers_a_version_that_says_what_it_learned_from(self): + from ami.jobs.models import Job + from ami.ml.models import Algorithm + + response = self._post({"result": self._result(), "dataset": self._dataset_echo()}) + + self.assertEqual(response.status_code, 200, response.data) + version = Algorithm.objects.get(key=response.data["algorithm"]) + self.assertEqual(version.training_info.occurrence_set_id, self.occurrence_set.pk) + self.assertEqual(version.training_info.occurrence_set_name, "Mine") + self.assertEqual(version.training_info.dataset_rows, 6) + self.assertEqual(version.training_info.metrics, {"top1": 0.8}) + self.assertEqual(version.training_info.previous_metrics, {"top1": 0.9}) + self.assertEqual(Job.objects.get(pk=self.job.pk).status, "SUCCESS") + + def test_an_echo_antenna_cannot_read_still_records_the_result(self): + """ + A service echoing an older shape loses the provenance, not the run. + + Refusing the whole callback would throw away a head that was trained successfully. + """ + from ami.ml.models import Algorithm + + response = self._post({"result": self._result(), "dataset": {"unexpected": "shape"}}) + + self.assertEqual(response.status_code, 200, response.data) + version = Algorithm.objects.get(key=response.data["algorithm"]) + self.assertIsNone(version.training_info.occurrence_set_id) + + def test_a_run_with_no_current_head_to_compare_against_is_recorded(self): + """ + A service reports a null incumbent rather than leaving the field out. + + Refusing that would throw away the result of a first-ever retrain, which is the + one run guaranteed to have nothing to compare against. + """ + from ami.ml.models import Algorithm + + response = self._post({"result": self._result(incumbent_metrics=None), "dataset": self._dataset_echo()}) + + self.assertEqual(response.status_code, 200, response.data) + version = Algorithm.objects.get(key=response.data["algorithm"]) + self.assertEqual(version.training_info.previous_metrics, {}) + + def test_a_result_antenna_cannot_read_is_refused(self): + """The result is what the new version is built from, so a broken one is a 400.""" + response = self._post({"result": {"labels": "not a list"}, "dataset": self._dataset_echo()}) + + self.assertEqual(response.status_code, 400) + + def test_a_second_callback_does_not_register_a_second_version(self): + from ami.ml.models import Algorithm + + self._post({"result": self._result(), "dataset": self._dataset_echo()}) + before = Algorithm.objects.count() + again = self._post({"result": self._result(), "dataset": self._dataset_echo()}) + + self.assertEqual(again.status_code, 200) + self.assertEqual(again.data["status"], "already recorded") + self.assertEqual(Algorithm.objects.count(), before) diff --git a/ami/ml/training/service.py b/ami/ml/training/service.py index da9866a9c..f684708f6 100644 --- a/ami/ml/training/service.py +++ b/ami/ml/training/service.py @@ -20,7 +20,7 @@ from django.urls import reverse from django.utils.text import slugify -from ami.ml.schemas import AlgorithmTrainingConfig +from ami.ml.schemas import AlgorithmTrainingConfig, TrainingRequest from ami.utils.requests import create_session, extract_error_message_from_response if TYPE_CHECKING: @@ -113,60 +113,52 @@ def absolute_media_url(url: str, base_url: str | None = None) -> str: return urljoin(base.rstrip("/") + "/", url.lstrip("/")) -def send_training_request( - job: "Job", service: "ProcessingService", algorithm: "Algorithm", dataset: dict -) -> dict | None: +def send_training_request(job: "Job", service: "ProcessingService", algorithm: "Algorithm", dataset: dict) -> None: """ Ask a processing service to retrain a head. - Returns the service's result if it answered inline, or None if it accepted the work and - will report back later. Raises if the service refused the request. + Returns once the service has accepted the work. The result is not read from the + response: it arrives at the job's result callback, which is the only way a run that + outlasts its request can report. A service that answers inline as well is simply + acknowledged. Raises if the service refused the request. """ endpoint = urljoin(service.endpoint_url.rstrip("/") + "/", "train") # The same merge the job validated before it built the dataset, so the service is told # the settings the training set was actually made under. Reading job.params again here # would let the two drift, and would send values nothing had checked. config = AlgorithmTrainingConfig.for_run(algorithm.training_config, job.params or {}) - payload: dict[str, typing.Any] = { - "dataset_url": absolute_media_url(dataset["url"]), - "algorithm_key": algorithm.key, - "job_id": job.pk, - "name": f"{algorithm.key}-job-{job.pk}", - "min_per_species": config.min_per_species, + request = TrainingRequest( + dataset_url=absolute_media_url(dataset["url"]), + algorithm_key=algorithm.key, + job_id=job.pk, + name=f"{algorithm.key}-job-{job.pk}", + min_per_species=config.min_per_species, # The fitting settings the service published, so an admin can tune them in Antenna # without redeploying the service. - "min_improvement": config.min_improvement, - "head_type": config.head_type, - "epochs": config.epochs, - "learning_rate": config.learning_rate, - "weight_decay": config.weight_decay, - } - - # Always sent: a service that finishes after the request times out reports back here - # instead, which is the only way a real training set can work. - payload["callback_url"] = callback_url_for(job) - payload["callback_token"] = make_callback_token(job) - # The head itself comes back here. A service that does not support the upload simply - # ignores this, and the version is registered without a stored copy as before. - payload["head_upload_url"] = head_upload_url_for(job) - # Optional for the service: without it the training stage simply stays at nought until - # the result lands, as it did before. - payload["progress_url"] = progress_url_for(job) + min_improvement=config.min_improvement, + head_type=config.head_type, + epochs=config.epochs, + learning_rate=config.learning_rate, + weight_decay=config.weight_decay, + # Always sent: a service that finishes after the request times out reports back + # here instead, which is the only way a real training set can work. + callback_url=callback_url_for(job), + callback_token=make_callback_token(job), + # The head itself comes back here. A service that does not support the upload + # ignores this, and the version is registered without a stored copy. + head_upload_url=head_upload_url_for(job), + # Without this the training stage simply stays at nought until the result lands. + progress_url=progress_url_for(job), + ) job.logger.info(f"Sending training request to {endpoint} for {algorithm.key}") session = create_session() - response = session.post(endpoint, json=payload, timeout=DISPATCH_TIMEOUT_SECONDS) + response = session.post(endpoint, json=request.dict(), timeout=DISPATCH_TIMEOUT_SECONDS) if not response.ok: message = extract_error_message_from_response(response) raise ValueError(f"The processing service refused the training request: {message}") - try: - return response.json() - except ValueError: - # Accepted, but nothing useful in the body. The service will report back. - return None - HEAD_DIRECTORY = "algorithms" From 2f55e09fce7d2553b51c722c4af138c09e5f1ac7 Mon Sep 17 00:00:00 2001 From: mohamedelabbas1996 Date: Fri, 9 Oct 2026 00:37:26 -0400 Subject: [PATCH 8/8] feat(jobs): say that a training job dispatches to an external service Every training job read as "internal", the mode for work the platform does itself, because Job.setup() decides the mode from the pipeline and a training job has none. It is the platform's own word for what this job has always done: hand the work to a service and wait to be called back. --- ami/jobs/models.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/ami/jobs/models.py b/ami/jobs/models.py index 10ebcca2a..c8ec4753f 100644 --- a/ami/jobs/models.py +++ b/ami/jobs/models.py @@ -1064,6 +1064,10 @@ def run(cls, job: "Job"): job.progress.add_stage("Preparing training set", cls.STAGE_PREPARE) job.progress.add_stage("Sending to processing service", cls.STAGE_DISPATCH) job.progress.add_stage("Training", cls.STAGE_TRAIN) + # The work goes to an external service that reports back on its own, which is what + # this mode means. Set here rather than in Job.setup(), which reads it off the + # pipeline a training job does not have. + job.dispatch_mode = JobDispatchMode.ASYNC_API job.update_status(JobState.STARTED) job.started_at = datetime.datetime.now() job.finished_at = None