diff --git a/ami/base/pagination.py b/ami/base/pagination.py index 9ebca7b21..e3b8e6308 100644 --- a/ami/base/pagination.py +++ b/ami/base/pagination.py @@ -29,3 +29,15 @@ def _get_project(self): if hasattr(view, "get_active_project"): return view.get_active_project() return None + + +class TrainingDataPagination(LimitOffsetPagination): + """ + Paging for embedding rows, which are far bigger than a normal API row. + + A 1024-dimension vector serialises to roughly 20 KB of JSON, so the platform default + of 10 is uselessly small and an unbounded limit would return hundreds of megabytes. + """ + + default_limit = 200 + max_limit = 1000 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..5434571a1 --- /dev/null +++ b/ami/jobs/migrations/0024_train_classifier_job_type.py @@ -0,0 +1,31 @@ +# Generated by Django 4.2.10 on 2026-09-02 13:35 + +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/migrations/0025_generate_embeddings_job_type.py b/ami/jobs/migrations/0025_generate_embeddings_job_type.py new file mode 100644 index 000000000..5f6962190 --- /dev/null +++ b/ami/jobs/migrations/0025_generate_embeddings_job_type.py @@ -0,0 +1,32 @@ +# Generated by Django 4.2.10 on 2026-09-03 16:54 + +from django.db import migrations, models + + +class Migration(migrations.Migration): + dependencies = [ + ("jobs", "0024_train_classifier_job_type"), + ] + + 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"), + ("generate_embeddings", "Generate embeddings"), + ], + default="unknown", + max_length=255, + verbose_name="Job Type", + ), + ), + ] diff --git a/ami/jobs/migrations/0026_evaluation.py b/ami/jobs/migrations/0026_evaluation.py new file mode 100644 index 000000000..93a4b7066 --- /dev/null +++ b/ami/jobs/migrations/0026_evaluation.py @@ -0,0 +1,33 @@ +# Generated by Django 4.2.10 on 2026-09-13 15:49 + +from django.db import migrations, models + + +class Migration(migrations.Migration): + dependencies = [ + ("jobs", "0025_generate_embeddings_job_type"), + ] + + 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"), + ("generate_embeddings", "Generate embeddings"), + ("evaluate_algorithm", "Evaluate algorithm"), + ], + default="unknown", + max_length=255, + verbose_name="Job Type", + ), + ), + ] diff --git a/ami/jobs/models.py b/ami/jobs/models.py index ff65f31f2..5ac917be7 100644 --- a/ami/jobs/models.py +++ b/ami/jobs/models.py @@ -444,6 +444,22 @@ class JobType: name: str key: str + # What a job of this type cannot run without. The API refuses to create one that is + # missing any of these, so a gap is a 400 when the job is made rather than a failure + # minutes later when it runs. ``required_params`` are keys inside ``Job.params``; + # ``required_fields`` are fields on the job itself. + required_fields: tuple[str, ...] = () + required_params: tuple[str, ...] = () + + # Whether a person can start one from the UI. The rest are created by the platform: + # a data sync, an export, a regroup. Nothing should offer those as a choice. + user_creatable: bool = False + + # How long a job of this type may go untouched before the stale-job check treats it as + # dead. A type that hands work to an external service and waits to be called back is + # silent on purpose, so silence is not evidence that it died; those override this. + stalled_after_minutes: int = 10 + # @TODO Consider adding custom vocabulary for job types to be used in the UI # verb: str = "Sync" # present_participle: str = "syncing" @@ -460,6 +476,7 @@ def run(cls, job: "Job"): class MLJob(JobType): name = "ML pipeline" key = "ml" + user_creatable = True @classmethod def run(cls, job: "Job"): @@ -562,7 +579,14 @@ def run(cls, job: "Job"): cls.process_images(job, images) @classmethod - def process_images(cls, job, images): + def process_images(cls, job, images, reprocess_all_images: bool | None = None): + """ + `reprocess_all_images` defaults to the project flag. An embedding job overrides it, + because it must always send the existing detections: it needs the vectors for crops + that are already detected, not a fresh run of the detector. + """ + if reprocess_all_images is None: + reprocess_all_images = job.project.feature_flags.reprocess_all_images image_count = len(images) # Keep track of sub-tasks for saving results, pair with batch number save_tasks: list[tuple[int, AsyncResult]] = [] @@ -585,7 +609,7 @@ def process_images(cls, job, images): images=chunk, job_id=job.pk, project_id=job.project.pk, - reprocess_all_images=job.project.feature_flags.reprocess_all_images, + reprocess_all_images=reprocess_all_images, ) job.logger.info(f"Processed image batch {i+1} in {time.time() - request_sent:.2f}s") except Exception as e: @@ -920,6 +944,524 @@ def run(cls, job: "Job"): job.save() +class GenerateEmbeddingsJob(JobType): + """ + Fill in the embeddings that a classifier head needs to be retrained. + + Runs the same inference pipeline as an ML job, but only over detections that a person + has verified and that have no embedding yet. The backbone is frozen, so an embedding + is computed once and reused by every later retrain. + + Existing detections are sent with the request, so the detector does not run again and + no new detections are created. Only the vectors are new. + """ + + name = "Generate embeddings" + key = "generate_embeddings" + required_fields = ("pipeline",) + user_creatable = True + + @classmethod + def run(cls, job: "Job"): + if not job.pipeline: + raise ValueError("A generate_embeddings job needs a pipeline to run the backbone.") + if not job.project.feature_flags.store_classification_embeddings: + raise ValueError( + f"Project '{job.project}' does not have store_classification_embeddings enabled, " + "so the embeddings this job produces would be thrown away." + ) + + algorithm = cls.target_algorithm(job) + job.progress.add_stage("Finding detections", "collect") + job.progress.add_stage("Computing embeddings", "process") + job.update_status(JobState.STARTED) + job.started_at = datetime.datetime.now() + job.finished_at = None + job.progress.update_stage("collect", status=JobState.STARTED, progress=0) + job.save() + + images = cls.images_needing_embeddings(job, algorithm) + job.progress.add_stage_param("collect", "Algorithm", algorithm.key) + job.progress.add_stage_param("collect", "Captures", len(images)) + job.progress.update_stage("collect", status=JobState.SUCCESS, progress=1, total_images=len(images)) + job.save() + + if not images: + job.logger.info(f"Every verified detection already has an embedding from {algorithm.key}. Nothing to do.") + job.progress.update_stage("process", status=JobState.SUCCESS, progress=1) + job.finished_at = datetime.datetime.now() + job.update_status(JobState.SUCCESS, save=True) + return + + job.logger.info(f"Computing {algorithm.key} embeddings for detections in {len(images)} capture(s)") + before = cls.embedding_count(algorithm) + + # Reuse the ML job's dispatch loop rather than a parallel copy of it: batching, + # progress and result saving are already handled there. + MLJob.process_images(job, images, reprocess_all_images=True) + + created = cls.embedding_count(algorithm) - before + job.logger.info(f"Stored {created} new embedding(s)") + job.progress.add_stage_param("process", "Embeddings stored", created) + job.save() + + @classmethod + def target_algorithm(cls, job: "Job"): + """The pipeline's trainable algorithm. That is the one whose embeddings are worth storing.""" + candidates = list(job.pipeline.algorithms.filter(trainable=True)) if job.pipeline else [] + key = (job.params or {}).get("algorithm_key") + if key: + match = next((a for a in candidates if a.key == key), None) + if not match: + raise ValueError( + f"Algorithm '{key}' is not a trainable algorithm on pipeline '{job.pipeline}'. " + f"Trainable algorithms: {[a.key for a in candidates] or 'none'}" + ) + return match + if not candidates: + raise ValueError( + f"Pipeline '{job.pipeline}' has no algorithm marked trainable, so there is no head " + "whose embeddings are worth storing." + ) + if len(candidates) > 1: + raise ValueError( + f"Pipeline '{job.pipeline}' has several trainable algorithms " + f"({[a.key for a in candidates]}). Pass algorithm_key in the job params to choose one." + ) + return candidates[0] + + @classmethod + def images_needing_embeddings(cls, job: "Job", algorithm) -> list[SourceImage]: + """ + Captures holding at least one verified detection with no embedding from this algorithm. + + Verified only, by decision: an embedding is ~2 KB and only a labelled crop can train + a head, so embedding everything would cost gigabytes to no purpose. + """ + from ami.main.models import Detection + from ami.ml import training_data + + detections = ( + Detection.objects.filter(occurrence_id__in=training_data.verified_occurrence_ids(job.project)) + .exclude(embeddings__algorithm=algorithm) + .order_by() + ) + image_ids = list(detections.values_list("source_image_id", flat=True).distinct()) + if job.limit: + image_ids = image_ids[: job.limit] + return list(SourceImage.objects.filter(pk__in=image_ids)) + + @classmethod + def embedding_count(cls, algorithm) -> int: + from ami.ml.models import DetectionEmbedding + + return DetectionEmbedding.objects.filter(algorithm=algorithm).count() + + +class EvaluateAlgorithmJob(JobType): + """ + Score an algorithm against a fixed set of occurrences people have verified. + + Reads predictions that already exist rather than running anything, so an algorithm that + has processed the set is scored in one pass with no GPU. Comparing two models means + scoring both on the same set, which is why the set's membership is fixed. + """ + + name = "Evaluate algorithm" + key = "evaluate_algorithm" + required_params = ("algorithm_key", "occurrence_set_id") + user_creatable = True + + STAGE_SCORE = "score" + + @classmethod + def run(cls, job: "Job"): + from ami.main.models import OccurrenceSet + from ami.ml import evaluation + from ami.ml.models import Algorithm + + params = job.params or {} + algorithm_key = params.get("algorithm_key") + occurrence_set_id = params.get("occurrence_set_id") + if not algorithm_key or not occurrence_set_id: + raise ValueError("An evaluate_algorithm job needs 'algorithm_key' and 'occurrence_set_id' in its params.") + + algorithm = Algorithm.objects.filter(key=algorithm_key).first() + if not algorithm: + raise ValueError(f"No algorithm with key '{algorithm_key}'.") + # Scoped to the project: the id comes from the job's params, which any member who + # can create a job may set. An unscoped lookup would score a model against another + # project's occurrences and, because an evaluation is stored once per algorithm and + # set, overwrite that project's own result. + 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.") + + job.progress.add_stage("Scoring", cls.STAGE_SCORE) + job.update_status(JobState.STARTED) + job.started_at = datetime.datetime.now() + job.finished_at = None + job.progress.update_stage(cls.STAGE_SCORE, status=JobState.STARTED, progress=0) + job.save() + + try: + result = evaluation.score(occurrence_set, algorithm) + except evaluation.NothingToScore as e: + # A missing step, not a bad model. Saying so beats recording an accuracy of zero. + job.logger.error(str(e)) + job.progress.update_stage(cls.STAGE_SCORE, status=JobState.FAILURE, progress=0) + job.finished_at = datetime.datetime.now() + job.result = {"error": str(e)} + job.update_status(JobState.FAILURE, save=True) + return + + stored = evaluation.save_evaluation(occurrence_set, algorithm, result, job=job) + + job.logger.info( + f"{algorithm.key} on '{occurrence_set.name}': " + f"{result['micro_accuracy']:.3f} overall, {result['macro_accuracy']:.3f} averaged over " + f"{result['species_scored']} species" + ) + # Said out loud because an accuracy over a handful of species reads the same as one + # over all of them. A score on 6 of 749 is a fact about this set, not about the model. + job.logger.info( + f"Scored {result['species_scored']} of the {result['species_in_set']} species in the set; " + f"{algorithm.key} can answer for {result['species_predictable']}." + ) + if result["occurrences_skipped"]: + job.logger.info( + f"{result['occurrences_skipped']} occurrence(s) were left out: their species is not one " + "this algorithm can predict, or it never classified them." + ) + job.progress.add_stage_param(cls.STAGE_SCORE, "Accuracy", round(result["micro_accuracy"], 3)) + job.progress.add_stage_param(cls.STAGE_SCORE, "Averaged over species", round(result["macro_accuracy"], 3)) + job.progress.add_stage_param(cls.STAGE_SCORE, "Occurrences scored", result["occurrences_scored"]) + job.progress.add_stage_param( + cls.STAGE_SCORE, "Species", f"{result['species_scored']} of {result['species_in_set']} in set" + ) + job.progress.add_stage_param(cls.STAGE_SCORE, "Species the model knows", result["species_predictable"]) + job.progress.update_stage(cls.STAGE_SCORE, status=JobState.SUCCESS, progress=1) + job.result = {"evaluation_id": stored.pk, "micro_accuracy": result["micro_accuracy"]} + job.finished_at = datetime.datetime.now() + job.update_status(JobState.SUCCESS, save=True) + job.save() + + +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",) + 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." + ) + + 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) + # The algorithm's training_config holds the defaults the service published; a job + # may override any of them for one run. + config = algorithm.training_config + try: + dataset = build_training_dataset( + project=job.project, + algorithm=algorithm, + min_per_species=params.get("min_per_species", config.min_per_species), + test_fraction=params.get("test_fraction", config.test_fraction), + split_salt=params.get("split_salt", config.split_salt), + taxa_list=taxa_list, + 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.get("taxa_list"): + job.logger.info( + f"Species list comes from taxa list '{meta['taxa_list']['name']}' " f"({len(meta['classes'])} species)" + ) + without_data = meta.get("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.get("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 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: + from ami.ml.models.training_set import attach_algorithm + + 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}") + # The occurrences were recorded when the set was built, before this version + # existed. Point them at it now. + linked = attach_algorithm(job=job, algorithm=new_version) + job.logger.info(f"Linked {linked} occurrence(s) to {new_version.key}") + + 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. + """ + from ami.ml.models import Algorithm, AlgorithmCategoryMap + from ami.ml.schemas import AlgorithmTrainingInfo + + result = payload.get("result") or {} + dataset = payload.get("dataset") or {} + 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"), + 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 [], + ), + ) + + class RegroupEventsJob(JobType): """ Regroup a deployment's captures into Events using the project's @@ -985,6 +1527,9 @@ def run(cls, job: "Job"): UnknownJobType, DataExportJob, PostProcessingJob, + TrainClassifierJob, + GenerateEmbeddingsJob, + EvaluateAlgorithmJob, ] @@ -1021,7 +1566,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/serializers.py b/ami/jobs/serializers.py index f53199e73..979ba062f 100644 --- a/ami/jobs/serializers.py +++ b/ami/jobs/serializers.py @@ -14,7 +14,16 @@ from ami.ml.schemas import PipelineProcessingTask, PipelineTaskResult, ProcessingServiceClientInfo from ami.ml.serializers import PipelineNestedSerializer -from .models import JOB_LOGS_DEFAULT_LIMIT, Job, JobProgress, MLJob, _legacy_logs_shape, serialize_job_logs +from .models import ( + JOB_LOGS_DEFAULT_LIMIT, + VALID_JOB_TYPES, + Job, + JobProgress, + MLJob, + _legacy_logs_shape, + get_job_type_by_key, + serialize_job_logs, +) from .schemas import QueuedTaskAcknowledgment @@ -52,9 +61,12 @@ class JobListSerializer(DefaultSerializer): progress = SchemaField(schema=JobProgress, read_only=True) logs = serializers.SerializerMethodField() job_type = JobTypeSerializer(read_only=True) - # All jobs created from the Jobs UI are ML jobs (datasync, etc. are created for the user) - # @TODO Remove this when the UI is updated pass a job type. This should be a required field. + # Defaulted rather than required: an API client written before job types were + # selectable still means an ML job. job_type_key = serializers.SlugField(write_only=True, default=MLJob.key) + # Free-form, and read by the job type that consumes it. A train_classifier job + # carries its algorithm_key here; an ML job ignores it entirely. + params = serializers.JSONField(required=False, allow_null=True) project_id = serializers.PrimaryKeyRelatedField( label="Project", @@ -129,6 +141,7 @@ class Meta: "logs", "job_type", "job_type_key", + "params", "data_export", "dispatch_mode", # "duration", @@ -148,6 +161,42 @@ class Meta: "dispatch_mode", ] + def validate_job_type_key(self, value): + job_type = get_job_type_by_key(value) + if not job_type: + known = sorted(t.key for t in VALID_JOB_TYPES) + raise serializers.ValidationError(f"Unknown job type '{value}'. Known types: {known}") + return value + + def validate(self, attrs): + """ + Refuse a job that could never run. + + Each job type says what it cannot run without. Checking here means a missing + pipeline or algorithm is reported while the person is still looking at the form, + rather than as a failed job minutes later. + """ + attrs = super().validate(attrs) + job_type = get_job_type_by_key(attrs.get("job_type_key", MLJob.key)) + if not job_type: + return attrs + + missing_fields = [name for name in job_type.required_fields if not attrs.get(name)] + if missing_fields: + raise serializers.ValidationError( + # Plural, so the message does not have to choose between "a" and "an": + # two of the four job type names start with a vowel sound. + {f"{name}_id": f"{job_type.name} jobs need a {name}." for name in missing_fields} + ) + + params = attrs.get("params") or {} + missing_params = [name for name in job_type.required_params if not params.get(name)] + if missing_params: + raise serializers.ValidationError( + {"params": f"{job_type.name} jobs need {', '.join(missing_params)} in their params."} + ) + return attrs + @extend_schema_field( { "type": "object", diff --git a/ami/jobs/tasks.py b/ami/jobs/tasks.py index b00213ebb..f0a9214f9 100644 --- a/ami/jobs/tasks.py +++ b/ami/jobs/tasks.py @@ -988,7 +988,7 @@ def check_stale_jobs(minutes: int | None = None, dry_run: bool = False) -> list[ from celery.result import AsyncResult from django.db import transaction - from ami.jobs.models import Job, JobDispatchMode, JobState + from ami.jobs.models import Job, JobDispatchMode, JobState, get_job_type_by_key if minutes is None: minutes = Job.STALLED_JOBS_MAX_MINUTES @@ -1014,6 +1014,17 @@ def check_stale_jobs(minutes: int | None = None, dry_run: bool = False) -> list[ # Another concurrent run already handled this job. continue + # Each job type says how long one of its jobs may go untouched. A type that + # hands work to an external service and waits to be called back writes nothing + # to its own row meanwhile, so silence is not evidence that it died. Candidates + # are gathered above at the default, then each is judged by its own deadline. + # Looked up rather than read off the job so an unrecognised job_type_key leaves + # that one job alone instead of raising and stopping the whole sweep. + job_type = get_job_type_by_key(job.job_type_key) + limit = job_type.stalled_after_minutes if job_type else minutes + if job.updated_at > datetime.datetime.now() - datetime.timedelta(minutes=limit): + continue + celery_state = None if job.task_id: try: diff --git a/ami/jobs/tests/test_jobs.py b/ami/jobs/tests/test_jobs.py index 00b7934a7..f57a8cbbf 100644 --- a/ami/jobs/tests/test_jobs.py +++ b/ami/jobs/tests/test_jobs.py @@ -10,6 +10,7 @@ from ami.base.serializers import reverse_with_params from ami.jobs.models import ( DataStorageSyncJob, + EvaluateAlgorithmJob, Job, JobDispatchMode, JobLog, @@ -18,9 +19,10 @@ MLJob, RegroupEventsJob, SourceImageCollectionPopulateJob, + TrainClassifierJob, ) -from ami.main.models import Deployment, Event, Project, SourceImage, SourceImageCollection -from ami.ml.models import Pipeline +from ami.main.models import Deployment, Event, OccurrenceSet, Project, SourceImage, SourceImageCollection, TaxaList +from ami.ml.models import Algorithm, Pipeline from ami.ml.models.processing_service import ProcessingService from ami.ml.orchestration.jobs import queue_images_to_nats from ami.tests.fixtures.main import create_captures @@ -328,6 +330,61 @@ def test_create_job(self): # progress = JobProgress(**data["progress"]) # self.assertEqual(progress.summary.status, JobState.CREATED) + def test_creating_a_job_of_an_unknown_type_is_refused(self): + jobs_create_url = reverse_with_params("api:job-list", params={"project_id": self.project.pk}) + self.client.force_authenticate(user=self.user) + + resp = self.client.post( + jobs_create_url, + {"project_id": self.project.pk, "name": "Nonsense", "delay": 0, "job_type_key": "not-a-job-type"}, + ) + + self.assertEqual(resp.status_code, 400) + self.assertIn("job_type_key", resp.json()) + + def test_a_job_missing_what_its_type_needs_is_refused(self): + """ + The gap is reported while the form is still open, rather than as a job that fails + minutes later for want of an algorithm. + """ + jobs_create_url = reverse_with_params("api:job-list", params={"project_id": self.project.pk}) + self.client.force_authenticate(user=self.user) + + resp = self.client.post( + jobs_create_url, + { + "project_id": self.project.pk, + "name": "Evaluate nothing in particular", + "delay": 0, + "job_type_key": EvaluateAlgorithmJob.key, + }, + format="json", + ) + + self.assertEqual(resp.status_code, 400) + self.assertIn("params", resp.json()) + + def test_a_job_carrying_what_its_type_needs_is_created(self): + jobs_create_url = reverse_with_params("api:job-list", params={"project_id": self.project.pk}) + self.client.force_authenticate(user=self.user) + + resp = self.client.post( + jobs_create_url, + { + "project_id": self.project.pk, + "name": "Evaluate a head", + "delay": 0, + "job_type_key": EvaluateAlgorithmJob.key, + "params": {"algorithm_key": "some-head", "occurrence_set_id": 1}, + }, + format="json", + ) + + self.assertEqual(resp.status_code, 201) + job = Job.objects.get(pk=resp.json()["id"]) + self.assertEqual(job.job_type_key, EvaluateAlgorithmJob.key) + self.assertEqual(job.params["algorithm_key"], "some-head") + def test_run_job(self): data = self._create_job("Test run job", start_now=False) @@ -1744,3 +1801,115 @@ def test_browsable_page_renders_number_input(self): html = response.content.decode() self.assertNotIn('