diff --git a/ami/base/pagination.py b/ami/base/pagination.py index 9ebca7b21..ffcbda445 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/base/permissions.py b/ami/base/permissions.py index 328aa8384..af2636c94 100644 --- a/ami/base/permissions.py +++ b/ami/base/permissions.py @@ -198,6 +198,35 @@ def has_permission(self, request, view): return super().has_permission(request, view) +class OccurrenceSetPermission(ObjectPermission): + """ + Permission for the occurrence-set route, where list and create have no object yet. + + Listing follows the project: anyone who may see the project may see which sets it has, + since only a set's name and size are exposed. Creating one is gated on the project's + ``create_occurrenceset`` permission. Everything else falls through to the object check, + which refuses a global set because it has no project to check against. + """ + + def has_permission(self, request, view): + from ami.main.models import Project + + if view.action == "list": + # The project id is required for this action, so reaching here means it was + # given and the caller may see it. + return view.get_active_project() is not None + + if view.action == "create": + # Read from the payload rather than the query string: the set names the project + # it will belong to, and that is the project whose permission must allow it. + project = Project.objects.filter(pk=request.data.get("project_id")).first() + if not project: + return False + return request.user.has_perm(Project.Permissions.CREATE_OCCURRENCE_SET, project) + + return super().has_permission(request, view) + + class UserMembershipPermission(ObjectPermission): """ Custom permission for UserProjectMembershipViewSet. 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 ff65f31f2..c8ec4753f 100644 --- a/ami/jobs/models.py +++ b/ami/jobs/models.py @@ -17,13 +17,17 @@ 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 logger = logging.getLogger(__name__) +if typing.TYPE_CHECKING: + from ami.ml.models import Algorithm + from ami.ml.schemas import TrainingResult + class JobDispatchMode(models.TextChoices): """ @@ -444,6 +448,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 +480,10 @@ def run(cls, job: "Job"): class MLJob(JobType): name = "ML pipeline" key = "ml" + # The one type a person starts by hand today. A pipeline is not declared required here: + # single-capture runs created from a session page carry one, but the API has always + # accepted a job without it and failing it at creation would change that. + user_creatable = True @classmethod def run(cls, job: "Job"): @@ -977,6 +1001,429 @@ 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",) + 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" + + # 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 + from ami.ml.models.processing_service import ProcessingService + from ami.ml.training 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.") + + # 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: + 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) + # 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 + 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") -> OccurrenceSet | None: + """ + The set of occurrences this run learns from, or None for every verified one. + + 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: + 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: + 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 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) + + send_training_request(job=job, service=service, algorithm=algorithm, dataset=dataset) + + # 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.") + 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.save(update_fields=["progress", "updated_at"]) + # 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: + """ + 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", + 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 None + + 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. + 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.warnings: + job.logger.warning(warning) + + 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.promote) + job.logger.info(result.reason or "Training finished.") + + 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}") + + 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() + return new_version + + @classmethod + 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. + + 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 + + # 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: + try: + 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.labels + if not labels: + job.logger.warning("The service returned no class list, so no new version was registered.") + return None + + 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() + 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 {metadata.rows if metadata else 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.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=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.rows.get("kept"), + dataset_classes=len(labels), + metrics=result.candidate_metrics, + previous_metrics=result.incumbent_metrics or {}, + parent_algorithm_key=parent.key, + job_id=job.pk, + warnings=result.warnings, + ), + ) + + VALID_JOB_TYPES = [ MLJob, SourceImageCollectionPopulateJob, @@ -985,6 +1432,7 @@ def run(cls, job: "Job"): UnknownJobType, DataExportJob, PostProcessingJob, + TrainClassifierJob, ] @@ -1021,7 +1469,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 206d3a5a5..ac521f39e 100644 --- a/ami/jobs/serializers.py +++ b/ami/jobs/serializers.py @@ -11,10 +11,18 @@ ) 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 JOB_LOGS_DEFAULT_LIMIT, Job, JobProgress, MLJob, _legacy_logs_shape, serialize_job_logs +from .models import ( + JOB_LOGS_DEFAULT_LIMIT, + Job, + JobProgress, + MLJob, + _legacy_logs_shape, + get_job_type_by_key, + serialize_job_logs, +) from .schemas import QueuedTaskAcknowledgment @@ -55,6 +63,10 @@ class JobListSerializer(DefaultSerializer): # 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. job_type_key = serializers.SlugField(write_only=True, default=MLJob.key) + # Settings a job type reads from its own params, such as which algorithm to retrain and + # what split to hold out. Kept as one field rather than a column per type: only the type + # knows what its params mean, and it is the type that validates them when it runs. + 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,40 @@ class Meta: "dispatch_mode", ] + def validate(self, attrs): + """ + Refuse a job its type could not run, while it is still being created. + + A missing setting is otherwise found minutes later by the worker, which leaves a + failed job in the list and a person wondering what they got wrong. Only creation is + checked: an existing job's type cannot be changed. + """ + attrs = super().validate(attrs) + if self.instance is not None: + return attrs + + job_type_key = attrs.get("job_type_key") or MLJob.key + job_type = get_job_type_by_key(job_type_key) + if not job_type: + raise serializers.ValidationError({"job_type_key": f"Unknown job type '{job_type_key}'."}) + # ``user_creatable`` is not enforced here: several types the platform creates as a + # side effect are also posted directly by existing clients. It says which types a + # job form should offer as a choice, which is the UI's question, not the API's. + + missing_fields = [name for name in job_type.required_fields if not attrs.get(name)] + if missing_fields: + raise serializers.ValidationError( + {f"{name}_id": f"A {job_type.name} job needs this." 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"A {job_type.name} job needs {', '.join(missing_params)} in its params."} + ) + return attrs + @extend_schema_field( { "type": "object", @@ -245,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/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/views.py b/ami/jobs/views.py index 09b314e4b..175ea70b4 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 @@ -28,6 +31,9 @@ MLJobResultsResponseSerializer, MLJobTasksRequestSerializer, MLJobTasksResponseSerializer, + TrainingProgressRequestSerializer, + TrainingResultRequestSerializer, + TrainingResultResponseSerializer, ) from ami.jobs.tasks import ( HEARTBEAT_THROTTLE_SECONDS, @@ -591,3 +597,156 @@ 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=TrainingResultRequestSerializer, + responses={200: TrainingResultResponseSerializer}, + ) + @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 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 + + 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 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", "job_id": job.pk, "algorithm": None}) + + serializer = TrainingResultRequestSerializer(data=request.data) + serializer.is_valid(raise_exception=True) + + 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", + "job_id": job.pk, + "algorithm": new_version.key if new_version else None, + } + ) + + @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"}) + + 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"], + ) + 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/main/api/serializers.py b/ami/main/api/serializers.py index 1ca588980..5611106ff 100644 --- a/ami/main/api/serializers.py +++ b/ami/main/api/serializers.py @@ -26,6 +26,7 @@ Event, Identification, Occurrence, + OccurrenceSet, Page, Project, ProjectSettingsMixin, @@ -1418,6 +1419,59 @@ def to_representation(self, instance): return {key: value for key, value in data.items() if value is not None} +class OccurrenceSetSerializer(DefaultSerializer): + """A fixed list of occurrences. ``occurrence_ids`` is accepted only when it is created.""" + + project_id = serializers.PrimaryKeyRelatedField(queryset=Project.objects.all(), write_only=True) + occurrence_ids = serializers.PrimaryKeyRelatedField( + queryset=Occurrence.objects.all(), many=True, write_only=True, allow_empty=False + ) + occurrences_count = serializers.SerializerMethodField() + + class Meta: + model = OccurrenceSet + fields = [ + "id", + "details", + "name", + "description", + "project_id", + "occurrence_ids", + "occurrences_count", + "created_at", + "updated_at", + ] + + def get_occurrences_count(self, obj) -> int: + return getattr(obj, "annotated_occurrences_count", None) or obj.occurrences.count() + + def validate(self, attrs): + # Scores recorded against a set were measured on these occurrences, so moving them + # would change what those numbers mean. Another set is the way to change the list. + if self.instance is not None and "occurrence_ids" in attrs: + raise serializers.ValidationError( + {"occurrence_ids": "A set's occurrences cannot be changed. Create another set instead."} + ) + + project = attrs.get("project_id") + occurrences = attrs.get("occurrence_ids") + if project and occurrences: + outside = sorted(o.pk for o in occurrences if o.project_id != project.pk) + if outside: + raise serializers.ValidationError( + {"occurrence_ids": f"Not occurrences in this project: {outside[:10]}"} + ) + return attrs + + def create(self, validated_data): + project = validated_data.pop("project_id") + occurrences = validated_data.pop("occurrence_ids") + occurrence_set = OccurrenceSet.objects.create(**validated_data) + occurrence_set.projects.add(project) + occurrence_set.occurrences.set(occurrences) + return occurrence_set + + class SourceImageCollectionSerializer(DefaultSerializer): source_images = serializers.SerializerMethodField() kwargs = SourceImageCollectionCommonKwargsSerializer(required=False, partial=True) diff --git a/ami/main/api/views.py b/ami/main/api/views.py index 65dde9f4e..b2805d94b 100644 --- a/ami/main/api/views.py +++ b/ami/main/api/views.py @@ -29,7 +29,12 @@ from ami.base.metadata import ResponseSchemaMetadata from ami.base.models import BaseQuerySet from ami.base.pagination import LimitOffsetPaginationWithPermissions -from ami.base.permissions import IsActiveStaffOrReadOnly, IsProjectMemberOrReadOnly, ObjectPermission +from ami.base.permissions import ( + IsActiveStaffOrReadOnly, + IsProjectMemberOrReadOnly, + ObjectPermission, + OccurrenceSetPermission, +) from ami.base.serializers import FilterParamsSerializer, SingleParamSerializer from ami.base.views import ProjectMixin from ami.main.api.schemas import limit_doc_param, project_id_doc_param @@ -50,6 +55,7 @@ Event, Identification, Occurrence, + OccurrenceSet, Page, Project, ProjectQuerySet, @@ -85,6 +91,7 @@ ModelAgreementSerializer, OccurrenceListSerializer, OccurrenceSerializer, + OccurrenceSetSerializer, PageListSerializer, PageSerializer, ProjectListSerializer, @@ -954,6 +961,74 @@ class ChoicesPagination(LimitOffsetPaginationWithPermissions): max_limit = 100 +class OccurrenceSetViewSet(DefaultViewSet, ProjectMixin): + """ + Fixed lists of occurrences, so the same data can be used again later. + + Membership is decided when a set is created and cannot be changed afterwards: there is + no endpoint to add or remove an occurrence. Anything that compares results over time + depends on the list standing still, and a set that grows quietly makes every number + recorded against it mean something different. + + A set belonging to no project is global and is offered to every project, but it cannot + be edited here, because it has no single project whose permissions would govern it. + """ + + queryset = OccurrenceSet.objects.all() + serializer_class = OccurrenceSetSerializer + permission_classes = [OccurrenceSetPermission] + # No put: a full replace would include the occurrences. + http_method_names = ["get", "post", "patch", "delete", "head", "options"] + ordering_fields = ["name", "created_at", "updated_at"] + search_fields = ["name"] + ordering = ["name"] + # Listing without a project would answer with every set on the platform, so it is + # required there. A detail route names one set, which carries its own project, and the + # create payload names the project it belongs to. + require_project = False + require_project_for_list = True + + def create(self, request, *args, **kwargs): + """ + Create without the base class's unsaved-instance permission check. + + ``DefaultViewSet.create`` builds ``Model(**validated_data)`` to check object + permissions before saving. A set's project and its occurrences arrive as payload + fields rather than columns, so that instance cannot be built. The same check runs + earlier instead, in ``OccurrenceSetPermission``, against the project the payload + names. + """ + serializer = self.get_serializer(data=request.data) + serializer.is_valid(raise_exception=True) + self.perform_create(serializer) + headers = self.get_success_headers(serializer.data) + return Response(serializer.data, status=status.HTTP_201_CREATED, headers=headers) + + @extend_schema(parameters=[project_id_doc_param], responses=OccurrenceSetSerializer(many=True)) + @action(detail=False, methods=["get"], name="choices") + def choices(self, request: Request) -> Response: + """ + Choices for the occurrence-set filter and pickers. + + Follows SourceImageCollectionViewSet.choices: most recently updated first and + enough of them that a dropdown never has to page. + """ + self.ordering_fields = ["id", "created_at", "updated_at", "name"] + queryset = self.filter_queryset(self.get_queryset()) + paginator = ChoicesPagination() + page = paginator.paginate_queryset(queryset, request, view=self) + serializer = self.get_serializer(page, many=True) + return paginator.get_paginated_response(serializer.data) + + def get_queryset(self) -> QuerySet["OccurrenceSet"]: + qs = super().get_queryset().annotate(annotated_occurrences_count=models.Count("occurrences")) + project = self.get_active_project() + if project is not None: + # A set with no project is global, so it is offered everywhere. + return qs.for_project(project) + return qs.visible_for_user(self.request.user) + + class SourceImageCollectionViewSet(DefaultViewSet, ProjectMixin): """ Endpoint for viewing capture sets or samples of captures. @@ -1496,6 +1571,9 @@ class OccurrenceFilterSet(FilterSet): """ detections__source_image = RelatedIdFilter() + # Named for what it means to a reader rather than for the reverse accessor, which is + # called evaluation_sets because scoring was the first thing to use one. + occurrence_set = RelatedIdFilter(field_name="evaluation_sets") class Meta: model = Occurrence diff --git a/ami/main/migrations/0098_occurrence_set.py b/ami/main/migrations/0098_occurrence_set.py new file mode 100644 index 000000000..26153b846 --- /dev/null +++ b/ami/main/migrations/0098_occurrence_set.py @@ -0,0 +1,102 @@ +# Generated by Django 4.2.10 on 2026-10-08 13:45 + +from django.db import migrations, models + + +class Migration(migrations.Migration): + dependencies = [ + ("main", "0097_detection_and_classification_job_indexes"), + ] + + operations = [ + migrations.AlterModelOptions( + name="project", + options={ + "ordering": ["-priority", "created_at"], + "permissions": [ + ("create_identification", "Can create identifications"), + ("update_identification", "Can update identifications"), + ("delete_identification", "Can delete identifications"), + ("create_job", "Can create a job"), + ("update_job", "Can update a job"), + ("run_ml_job", "Can run/retry/cancel ML jobs"), + ("run_populate_captures_collection_job", "Can run/retry/cancel Populate Collection jobs"), + ("run_data_storage_sync_job", "Can run/retry/cancel Data Storage Sync jobs"), + ("run_regroup_events_job", "Can run/retry/cancel Regroup Events jobs"), + ("run_data_export_job", "Can run/retry/cancel Data Export jobs"), + ("run_single_image_ml_job", "Can process a single capture"), + ("run_post_processing_job", "Can run/retry/cancel Post-Processing jobs"), + ("delete_job", "Can delete a job"), + ("create_deployment", "Can create a deployment"), + ("delete_deployment", "Can delete a deployment"), + ("update_deployment", "Can update a deployment"), + ("sync_deployment", "Can sync images to a deployment"), + ("regroup_sessions_deployment", "Can regroup deployment captures into sessions"), + ("create_occurrenceset", "Can create an occurrence set"), + ("update_occurrenceset", "Can rename or describe an occurrence set"), + ("delete_occurrenceset", "Can delete an occurrence set"), + ("create_sourceimagecollection", "Can create a collection"), + ("update_sourceimagecollection", "Can update a collection"), + ("delete_sourceimagecollection", "Can delete a collection"), + ("populate_sourceimagecollection", "Can populate a collection"), + ("create_sourceimage", "Can create a source image"), + ("update_sourceimage", "Can update a source image"), + ("delete_sourceimage", "Can delete a source image"), + ("star_sourceimage", "Can star a source image"), + ("create_sourceimageupload", "Can create a source image upload"), + ("update_sourceimageupload", "Can update a source image upload"), + ("delete_sourceimageupload", "Can delete a source image upload"), + ("create_s3storagesource", "Can create storage"), + ("delete_s3storagesource", "Can delete storage"), + ("update_s3storagesource", "Can update storage"), + ("test_s3storagesource", "Can test storage connection"), + ("create_site", "Can create a site"), + ("delete_site", "Can delete a site"), + ("update_site", "Can update a site"), + ("create_device", "Can create a device"), + ("delete_device", "Can delete a device"), + ("update_device", "Can update a device"), + ("view_userprojectmembership", "Can view project members"), + ("create_userprojectmembership", "Can add a user to the project"), + ("update_userprojectmembership", "Can update a user's project membership and role in the project"), + ("delete_userprojectmembership", "Can remove a user from the project"), + ("create_dataexport", "Can create a data export"), + ("update_dataexport", "Can update a data export"), + ("delete_dataexport", "Can delete a data export"), + ("create_projectpipelineconfig", "Can register pipelines for the project"), + ("update_projectpipelineconfig", "Can update pipeline configurations"), + ("delete_projectpipelineconfig", "Can remove pipelines from the project"), + ("create_taxalist", "Can create a taxa list"), + ("update_taxalist", "Can update a taxa list"), + ("delete_taxalist", "Can delete a taxa list"), + ("view_private_data", "Can view private data"), + ], + }, + ), + migrations.CreateModel( + name="OccurrenceSet", + fields=[ + ("id", models.BigAutoField(auto_created=True, primary_key=True, serialize=False, verbose_name="ID")), + ("created_at", models.DateTimeField(auto_now_add=True)), + ("updated_at", models.DateTimeField(auto_now=True)), + ("name", models.CharField(max_length=255)), + ("description", models.TextField(blank=True)), + ( + "occurrences", + models.ManyToManyField(blank=True, related_name="evaluation_sets", to="main.occurrence"), + ), + ( + "projects", + models.ManyToManyField( + blank=True, + help_text="Projects this set belongs to. A set with none is available everywhere.", + related_name="occurrence_sets", + to="main.project", + ), + ), + ], + options={ + "ordering": ["name"], + }, + ), + ] diff --git a/ami/main/migrations/0099_grant_occurrence_set_permissions.py b/ami/main/migrations/0099_grant_occurrence_set_permissions.py new file mode 100644 index 000000000..f69e2fcda --- /dev/null +++ b/ami/main/migrations/0099_grant_occurrence_set_permissions.py @@ -0,0 +1,98 @@ +""" +Grant the new occurrence-set permissions to the roles that already curate a project's data. + +``create_roles_for_project`` gives new projects their role permissions, and a +``post_migrate`` signal re-syncs every project on migrate, so this backfill is belt and +braces. It is kept so the grant is explicit rather than a side effect of a signal, and so +the reverse is written down. + +Guardian reads **object-level** rows: ``get_perms(user, project)`` and +``user.has_perm("create_occurrenceset", project)`` both look at +``GroupObjectPermission``. Adding the permission to ``group.permissions`` alone would not +show up there, so each project gets a row per role group. + +``ProjectManager`` inherits ``MLDataManager``'s permissions, so both groups are covered. +""" + +from django.db import migrations +from django.db.models import Q + +PERMISSIONS = ("create_occurrenceset", "update_occurrenceset", "delete_occurrenceset") +ROLE_GROUP_SUFFIXES = ("_MLDataManager", "_ProjectManager") + + +def _permissions(apps): + Permission = apps.get_model("auth", "Permission") + ContentType = apps.get_model("contenttypes", "ContentType") + try: + project_ct = ContentType.objects.get(app_label="main", model="project") + except ContentType.DoesNotExist: + return [], None + found = list(Permission.objects.filter(codename__in=PERMISSIONS, content_type=project_ct)) + return found, project_ct + + +def _project_pk_from_group(group): + # Group names are "{project_pk}_{project_name}_{RoleName}"; the pk is the immutable + # leading segment, since a project name may itself contain underscores. + try: + return int(group.name.split("_", 1)[0]) + except (ValueError, IndexError): + return None + + +def _role_groups(apps): + Group = apps.get_model("auth", "Group") + query = Q() + for suffix in ROLE_GROUP_SUFFIXES: + query |= Q(name__endswith=suffix) + return Group.objects.filter(query) + + +def grant(apps, schema_editor): + GroupObjectPermission = apps.get_model("guardian", "GroupObjectPermission") + permissions, project_ct = _permissions(apps) + if not permissions: + return + + for group in _role_groups(apps): + project_pk = _project_pk_from_group(group) + if project_pk is None: + continue + for permission in permissions: + group.permissions.add(permission) + GroupObjectPermission.objects.get_or_create( + permission=permission, + content_type=project_ct, + object_pk=str(project_pk), + group=group, + ) + + +def revoke(apps, schema_editor): + GroupObjectPermission = apps.get_model("guardian", "GroupObjectPermission") + permissions, project_ct = _permissions(apps) + if not permissions: + return + + for group in _role_groups(apps): + project_pk = _project_pk_from_group(group) + if project_pk is None: + continue + for permission in permissions: + group.permissions.remove(permission) + GroupObjectPermission.objects.filter( + permission=permission, + content_type=project_ct, + object_pk=str(project_pk), + group=group, + ).delete() + + +class Migration(migrations.Migration): + dependencies = [ + ("main", "0098_occurrence_set"), + ("guardian", "0002_generic_permissions_index"), + ] + + operations = [migrations.RunPython(grant, revoke)] diff --git a/ami/main/migrations/0100_train_classifier_job_permission.py b/ami/main/migrations/0100_train_classifier_job_permission.py new file mode 100644 index 000000000..ccb070eb1 --- /dev/null +++ b/ami/main/migrations/0100_train_classifier_job_permission.py @@ -0,0 +1,77 @@ +# Generated by Django 4.2.10 on 2026-10-08 16:04 + +from django.db import migrations + + +class Migration(migrations.Migration): + dependencies = [ + ("main", "0099_grant_occurrence_set_permissions"), + ] + + operations = [ + migrations.AlterModelOptions( + name="project", + options={ + "ordering": ["-priority", "created_at"], + "permissions": [ + ("create_identification", "Can create identifications"), + ("update_identification", "Can update identifications"), + ("delete_identification", "Can delete identifications"), + ("create_job", "Can create a job"), + ("update_job", "Can update a job"), + ("run_ml_job", "Can run/retry/cancel ML jobs"), + ("run_populate_captures_collection_job", "Can run/retry/cancel Populate Collection jobs"), + ("run_data_storage_sync_job", "Can run/retry/cancel Data Storage Sync jobs"), + ("run_regroup_events_job", "Can run/retry/cancel Regroup Events jobs"), + ("run_data_export_job", "Can run/retry/cancel Data Export jobs"), + ("run_single_image_ml_job", "Can process a single capture"), + ("run_post_processing_job", "Can run/retry/cancel Post-Processing jobs"), + ("run_train_classifier_job", "Can run/retry/cancel Train Classifier jobs"), + ("delete_job", "Can delete a job"), + ("create_deployment", "Can create a deployment"), + ("delete_deployment", "Can delete a deployment"), + ("update_deployment", "Can update a deployment"), + ("sync_deployment", "Can sync images to a deployment"), + ("regroup_sessions_deployment", "Can regroup deployment captures into sessions"), + ("create_occurrenceset", "Can create an occurrence set"), + ("update_occurrenceset", "Can rename or describe an occurrence set"), + ("delete_occurrenceset", "Can delete an occurrence set"), + ("create_sourceimagecollection", "Can create a collection"), + ("update_sourceimagecollection", "Can update a collection"), + ("delete_sourceimagecollection", "Can delete a collection"), + ("populate_sourceimagecollection", "Can populate a collection"), + ("create_sourceimage", "Can create a source image"), + ("update_sourceimage", "Can update a source image"), + ("delete_sourceimage", "Can delete a source image"), + ("star_sourceimage", "Can star a source image"), + ("create_sourceimageupload", "Can create a source image upload"), + ("update_sourceimageupload", "Can update a source image upload"), + ("delete_sourceimageupload", "Can delete a source image upload"), + ("create_s3storagesource", "Can create storage"), + ("delete_s3storagesource", "Can delete storage"), + ("update_s3storagesource", "Can update storage"), + ("test_s3storagesource", "Can test storage connection"), + ("create_site", "Can create a site"), + ("delete_site", "Can delete a site"), + ("update_site", "Can update a site"), + ("create_device", "Can create a device"), + ("delete_device", "Can delete a device"), + ("update_device", "Can update a device"), + ("view_userprojectmembership", "Can view project members"), + ("create_userprojectmembership", "Can add a user to the project"), + ("update_userprojectmembership", "Can update a user's project membership and role in the project"), + ("delete_userprojectmembership", "Can remove a user from the project"), + ("create_dataexport", "Can create a data export"), + ("update_dataexport", "Can update a data export"), + ("delete_dataexport", "Can delete a data export"), + ("create_projectpipelineconfig", "Can register pipelines for the project"), + ("update_projectpipelineconfig", "Can update pipeline configurations"), + ("delete_projectpipelineconfig", "Can remove pipelines from the project"), + ("create_taxalist", "Can create a taxa list"), + ("update_taxalist", "Can update a taxa list"), + ("delete_taxalist", "Can delete a taxa list"), + ("view_private_data", "Can view private data"), + ], + }, + ), + ] diff --git a/ami/main/migrations/0101_grant_train_classifier_job_permission.py b/ami/main/migrations/0101_grant_train_classifier_job_permission.py new file mode 100644 index 000000000..19c5cb8ec --- /dev/null +++ b/ami/main/migrations/0101_grant_train_classifier_job_permission.py @@ -0,0 +1,94 @@ +""" +Grant ``run_train_classifier_job`` to the role that already runs the other ML jobs. + +``create_roles_for_project`` gives new projects their role permissions, and a +``post_migrate`` signal re-syncs every project on migrate, so this backfill is belt and +braces. It is kept so the grant is explicit rather than a side effect of a signal. + +Guardian reads object-level rows, so each project gets a ``GroupObjectPermission``; adding +to ``group.permissions`` alone would not appear in ``get_perms(user, project)``. Mirrors +``0099_grant_occurrence_set_permissions``. +""" + +from django.db import migrations +from django.db.models import Q + +PERMISSION = "run_train_classifier_job" +ROLE_GROUP_SUFFIXES = ("_MLDataManager", "_ProjectManager") + + +def _permission(apps): + Permission = apps.get_model("auth", "Permission") + ContentType = apps.get_model("contenttypes", "ContentType") + try: + project_ct = ContentType.objects.get(app_label="main", model="project") + except ContentType.DoesNotExist: + return None, None + try: + return Permission.objects.get(codename=PERMISSION, content_type=project_ct), project_ct + except Permission.DoesNotExist: + return None, None + + +def _project_pk_from_group(group): + # Group names are "{project_pk}_{project_name}_{RoleName}"; the pk is the immutable + # leading segment, since a project name may itself contain underscores. + try: + return int(group.name.split("_", 1)[0]) + except (ValueError, IndexError): + return None + + +def _role_groups(apps): + Group = apps.get_model("auth", "Group") + query = Q() + for suffix in ROLE_GROUP_SUFFIXES: + query |= Q(name__endswith=suffix) + return Group.objects.filter(query) + + +def grant(apps, schema_editor): + GroupObjectPermission = apps.get_model("guardian", "GroupObjectPermission") + permission, project_ct = _permission(apps) + if permission is None: + return + + for group in _role_groups(apps): + project_pk = _project_pk_from_group(group) + if project_pk is None: + continue + group.permissions.add(permission) + GroupObjectPermission.objects.get_or_create( + permission=permission, + content_type=project_ct, + object_pk=str(project_pk), + group=group, + ) + + +def revoke(apps, schema_editor): + GroupObjectPermission = apps.get_model("guardian", "GroupObjectPermission") + permission, project_ct = _permission(apps) + if permission is None: + return + + for group in _role_groups(apps): + project_pk = _project_pk_from_group(group) + if project_pk is None: + continue + group.permissions.remove(permission) + GroupObjectPermission.objects.filter( + permission=permission, + content_type=project_ct, + object_pk=str(project_pk), + group=group, + ).delete() + + +class Migration(migrations.Migration): + dependencies = [ + ("main", "0100_train_classifier_job_permission"), + ("guardian", "0002_generic_permissions_index"), + ] + + operations = [migrations.RunPython(grant, revoke)] diff --git a/ami/main/migrations/0102_project_default_taxa_list.py b/ami/main/migrations/0102_project_default_taxa_list.py new file mode 100644 index 000000000..5634621b9 --- /dev/null +++ b/ami/main/migrations/0102_project_default_taxa_list.py @@ -0,0 +1,25 @@ +# Generated by Django 4.2.10 on 2026-10-08 16:11 + +from django.db import migrations, models +import django.db.models.deletion + + +class Migration(migrations.Migration): + dependencies = [ + ("main", "0101_grant_train_classifier_job_permission"), + ] + + operations = [ + migrations.AddField( + model_name="project", + name="default_taxa_list", + field=models.ForeignKey( + blank=True, + help_text="The species this project expects to see. Used as the class list when retraining a classifier head, so the head covers the region rather than only the species someone happened to verify.", + null=True, + on_delete=django.db.models.deletion.SET_NULL, + related_name="default_for_projects", + to="main.taxalist", + ), + ), + ] diff --git a/ami/main/models.py b/ami/main/models.py index 41323cb3d..616800277 100644 --- a/ami/main/models.py +++ b/ami/main/models.py @@ -429,6 +429,19 @@ def check_custom_permission(self, user, action: str) -> bool: # Fall back to default permission checking for other actions return super().check_custom_permission(user, action) + default_taxa_list = models.ForeignKey( + "TaxaList", + on_delete=models.SET_NULL, + null=True, + blank=True, + related_name="default_for_projects", + help_text=( + "The species this project expects to see. Used as the class list when retraining " + "a classifier head, so the head covers the region rather than only the species " + "someone happened to verify." + ), + ) + class Permissions: """CRUD Permission names follow the convention: `create_`, `update_`, `delete_`, `view_`""" @@ -454,6 +467,7 @@ class Permissions: RUN_REGROUP_EVENTS_JOB = "run_regroup_events_job" RUN_DATA_EXPORT_JOB = "run_data_export_job" RUN_POST_PROCESSING_JOB = "run_post_processing_job" + RUN_TRAIN_CLASSIFIER_JOB = "run_train_classifier_job" DELETE_JOB = "delete_job" # Deployment permissions @@ -463,6 +477,11 @@ class Permissions: SYNC_DEPLOYMENT = "sync_deployment" REGROUP_SESSIONS_DEPLOYMENT = "regroup_sessions_deployment" + # Occurrence set permissions + CREATE_OCCURRENCE_SET = "create_occurrenceset" + UPDATE_OCCURRENCE_SET = "update_occurrenceset" + DELETE_OCCURRENCE_SET = "delete_occurrenceset" + # Collection permissions CREATE_COLLECTION = "create_sourceimagecollection" UPDATE_COLLECTION = "update_sourceimagecollection" @@ -537,6 +556,7 @@ class Meta: ("run_data_export_job", "Can run/retry/cancel Data Export jobs"), ("run_single_image_ml_job", "Can process a single capture"), ("run_post_processing_job", "Can run/retry/cancel Post-Processing jobs"), + ("run_train_classifier_job", "Can run/retry/cancel Train Classifier jobs"), ("delete_job", "Can delete a job"), # Deployment permissions ("create_deployment", "Can create a deployment"), @@ -545,6 +565,9 @@ class Meta: ("sync_deployment", "Can sync images to a deployment"), ("regroup_sessions_deployment", "Can regroup deployment captures into sessions"), # Collection permissions + ("create_occurrenceset", "Can create an occurrence set"), + ("update_occurrenceset", "Can rename or describe an occurrence set"), + ("delete_occurrenceset", "Can delete an occurrence set"), ("create_sourceimagecollection", "Can create a collection"), ("update_sourceimagecollection", "Can update a collection"), ("delete_sourceimagecollection", "Can delete a collection"), @@ -4947,6 +4970,74 @@ def html(self) -> str: ] +class OccurrenceSetQuerySet(BaseQuerySet): + def for_project(self, project) -> models.QuerySet: + """Sets this project can use: its own, plus any that belong to no project.""" + return self.filter(models.Q(projects=project) | models.Q(projects__isnull=True)).distinct() + + def visible_for_user(self, user) -> models.QuerySet: + """ + Global sets stay visible; project sets follow their project. + + The inherited filter keeps a row only if it reaches a non-draft project, which a + set belonging to no project never does. Without this a global set is hidden from + everyone but a superuser, which is the opposite of what global means. Only a set's + name and size are exposed, never the occurrences inside it. + """ + visible = super().visible_for_user(user) + return self.filter(models.Q(pk__in=visible.values("pk")) | models.Q(projects__isnull=True)).distinct() + + +@final +class OccurrenceSet(BaseModel): + """ + A fixed list of occurrences, kept so the same ones can be used again later. + + Membership is stored rather than described by a filter, because a filter answers + differently as data arrives. Anything that compares results over time needs the list to + stand still: scoring two classifiers is only meaningful if both saw the same + occurrences. + + A set with no projects is global, which is how one set is shared across the platform; + that follows how TaxaList already treats a list with no project. + """ + + name = models.CharField(max_length=255) + description = models.TextField(blank=True) + occurrences = models.ManyToManyField("main.Occurrence", related_name="evaluation_sets", blank=True) + projects = models.ManyToManyField( + "main.Project", + related_name="occurrence_sets", + blank=True, + help_text="Projects this set belongs to. A set with none is available everywhere.", + ) + + objects = OccurrenceSetQuerySet.as_manager() + + class Meta: + ordering = ["name"] + + def __str__(self) -> str: + return f"{self.name} ({self.occurrences.count()} occurrences)" + + @property + def is_global(self) -> bool: + return not self.projects.exists() + + def get_project(self): + """ + The project whose permissions govern this set. + + BaseModel returns None for a many-to-many project relation, and + ``check_permission`` refuses every action without a project, so a global set cannot + be renamed or deleted through the API at all. That is the behaviour we want: a set + shared across projects changing under one of them would quietly change what every + other project's stored scores were measured on. Global sets stay a deliberate, + out-of-band act. + """ + return self.projects.first() + + class SourceImageCollectionQuerySet(BaseQuerySet): def with_source_images_count(self): return self.annotate( diff --git a/ami/main/test_occurrence_sets.py b/ami/main/test_occurrence_sets.py new file mode 100644 index 000000000..ad84c7c58 --- /dev/null +++ b/ami/main/test_occurrence_sets.py @@ -0,0 +1,254 @@ +""" +Occurrence sets: a fixed list of occurrences, decided once. + +The point of the model is that the list does not move. Anything comparing results over +time — scoring two classifiers, re-running an export — is only meaningful if both runs saw +the same occurrences. So these tests are mostly about what the API refuses. +""" + +import datetime +import uuid + +from rest_framework import status +from rest_framework.test import APITestCase + +from ami.main.models import ( + Classification, + Deployment, + Detection, + Event, + Occurrence, + OccurrenceSet, + Project, + S3StorageSource, + SourceImage, + Taxon, + TaxonRank, +) +from ami.users.models import User +from ami.users.roles import BasicMember, MLDataManager + + +class OccurrenceSetFixture(APITestCase): + """ + Occurrences are created directly rather than through the capture fixtures. + + These tests are about what the endpoint accepts and refuses, and need nothing from a + detection. Going through create_captures/create_occurrences pulls in event grouping and + a storage read for image dimensions, neither of which works in a test, and the failures + land at random on whichever test ran first. + """ + + def setUp(self) -> None: + self.owner = User.objects.create_user(email="owner@example.com", password="testpass123") + self.curator = User.objects.create_user(email="curator@example.com", password="testpass123") + self.member = User.objects.create_user(email="member@example.com", password="testpass123") + + self.project = Project.objects.create(name="Set Project", owner=self.owner) + MLDataManager.assign_user(self.curator, self.project) + BasicMember.assign_user(self.member, self.project) + + self.occurrences = [self.make_occurrence(self.project) for _ in range(3)] + self.url = "/api/v2/occurrences/sets/" + + @staticmethod + def make_occurrence(project: Project) -> Occurrence: + """ + One occurrence the list endpoint will actually return. + + Built here rather than through the capture fixtures, which time their images from + the current clock and group them into events by reading image dimensions from + storage. Neither works in a test, and the failures land at random on whichever test + ran first. The list only shows occurrences with a real detection and a + determination, so both are set. + """ + storage = S3StorageSource.objects.create(name=f"source {uuid.uuid4().hex[:6]}", project=project, bucket="test") + deployment = Deployment.objects.create( + name=f"station {uuid.uuid4().hex[:6]}", project=project, data_source=storage + ) + night = datetime.datetime(2026, 6, 1, 22, 0, 0, tzinfo=datetime.timezone.utc) + event = Event.objects.create(deployment=deployment, project=project, start=night) + image = SourceImage.objects.create( + deployment=deployment, + project=project, + event=event, + timestamp=night, + path=f"test/{uuid.uuid4().hex[:8]}.jpg", + ) + taxon = Taxon.objects.create(name=f"Testus {uuid.uuid4().hex[:6]}", rank=TaxonRank.SPECIES.name) + taxon.projects.add(project) + detection = Detection.objects.create(source_image=image, timestamp=night, bbox=[0.1, 0.1, 0.2, 0.2]) + Classification.objects.create(detection=detection, taxon=taxon, score=0.9, timestamp=night) + occurrence = detection.associate_new_occurrence() + occurrence.refresh_from_db() + return occurrence + + def _create(self, user, **overrides): + payload = { + "name": "Blind set", + "description": "", + "project_id": self.project.pk, + "occurrence_ids": [o.pk for o in self.occurrences], + } + payload.update(overrides) + self.client.force_authenticate(user=user) + return self.client.post(self.url, payload, format="json") + + +class TestCreatingASet(OccurrenceSetFixture): + def test_a_curator_can_create_one(self): + response = self._create(self.curator) + + self.assertEqual(response.status_code, status.HTTP_201_CREATED, response.data) + created = OccurrenceSet.objects.get(pk=response.data["id"]) + self.assertEqual(created.occurrences.count(), 3) + self.assertEqual(list(created.projects.all()), [self.project]) + + def test_a_member_without_the_permission_cannot(self): + response = self._create(self.member) + + self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN) + self.assertFalse(OccurrenceSet.objects.exists()) + + def test_an_anonymous_caller_cannot(self): + response = self._create(None) + + self.assertIn(response.status_code, (status.HTTP_401_UNAUTHORIZED, status.HTTP_403_FORBIDDEN)) + self.assertFalse(OccurrenceSet.objects.exists()) + + def test_an_occurrence_from_another_project_is_refused(self): + """Otherwise a set quietly reaches into data its project cannot see.""" + other = Project.objects.create(name="Other Project", owner=self.owner) + outsider = self.make_occurrence(other) + + response = self._create(self.curator, occurrence_ids=[self.occurrences[0].pk, outsider.pk]) + + self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST) + self.assertIn("occurrence_ids", response.data) + self.assertFalse(OccurrenceSet.objects.exists()) + + def test_an_empty_set_is_refused(self): + response = self._create(self.curator, occurrence_ids=[]) + + self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST) + + +class TestMembershipDoesNotMove(OccurrenceSetFixture): + def setUp(self) -> None: + super().setUp() + self.set = OccurrenceSet.objects.create(name="Blind set") + self.set.projects.add(self.project) + self.set.occurrences.add(*self.occurrences) + self.detail_url = f"{self.url}{self.set.pk}/" + self.client.force_authenticate(user=self.curator) + + def test_the_occurrences_cannot_be_changed(self): + response = self.client.patch(self.detail_url, {"occurrence_ids": [self.occurrences[0].pk]}, format="json") + + self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST) + self.assertEqual(self.set.occurrences.count(), 3) + + def test_the_name_and_description_can_be_changed(self): + response = self.client.patch(self.detail_url, {"name": "Renamed", "description": "why"}, format="json") + + self.assertEqual(response.status_code, status.HTTP_200_OK, response.data) + self.set.refresh_from_db() + self.assertEqual(self.set.name, "Renamed") + + def test_there_is_no_put(self): + """A full replace would carry the occurrences with it.""" + response = self.client.put(self.detail_url, {"name": "Replaced"}, format="json") + + self.assertEqual(response.status_code, status.HTTP_405_METHOD_NOT_ALLOWED) + + def test_a_curator_can_delete_one(self): + response = self.client.delete(self.detail_url) + + self.assertEqual(response.status_code, status.HTTP_204_NO_CONTENT) + self.assertFalse(OccurrenceSet.objects.filter(pk=self.set.pk).exists()) + + +class TestFilteringOccurrencesBySet(OccurrenceSetFixture): + """The occurrence list can be narrowed to one set, which is how a set is reviewed.""" + + def setUp(self) -> None: + super().setUp() + self.set = OccurrenceSet.objects.create(name="Blind set") + self.set.projects.add(self.project) + self.set.occurrences.add(self.occurrences[0], self.occurrences[1]) + self.client.force_authenticate(user=self.curator) + + def _ids(self, **params): + # apply_defaults=false: the list hides occurrences without a determination or a + # score by default, and these fixtures are bare rows made to exercise the filter. + response = self.client.get( + "/api/v2/occurrences/", + {"project_id": self.project.pk, "apply_defaults": "false", **params}, + ) + self.assertEqual(response.status_code, status.HTTP_200_OK, response.data) + return {row["id"] for row in response.data["results"]} + + def test_only_the_sets_occurrences_come_back(self): + self.assertEqual( + self._ids(occurrence_set=self.set.pk), + {self.occurrences[0].pk, self.occurrences[1].pk}, + ) + + def test_without_the_filter_every_occurrence_comes_back(self): + self.assertEqual(self._ids(), {o.pk for o in self.occurrences}) + + +class TestGlobalSets(OccurrenceSetFixture): + def setUp(self) -> None: + super().setUp() + self.global_set = OccurrenceSet.objects.create(name="Platform-wide set") + self.client.force_authenticate(user=self.curator) + + def _names(self, response): + return [row["name"] for row in response.data["results"]] + + def test_a_global_set_is_offered_to_every_project(self): + response = self.client.get(self.url, {"project_id": self.project.pk}) + + self.assertEqual(response.status_code, status.HTTP_200_OK) + self.assertIn(self.global_set.name, self._names(response)) + + def test_a_global_set_cannot_be_edited_through_the_api(self): + """It has no single project whose permissions would govern it.""" + response = self.client.patch(f"{self.url}{self.global_set.pk}/", {"name": "Mine now"}, format="json") + + self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN) + self.global_set.refresh_from_db() + self.assertEqual(self.global_set.name, "Platform-wide set") + + +class TestListing(OccurrenceSetFixture): + def setUp(self) -> None: + super().setUp() + self.set = OccurrenceSet.objects.create(name="Blind set") + self.set.projects.add(self.project) + self.set.occurrences.add(*self.occurrences) + + def test_a_project_is_required(self): + """Without one the endpoint would list every set on the platform.""" + self.client.force_authenticate(user=self.curator) + response = self.client.get(self.url) + + self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST) + + def test_the_size_is_reported_without_the_occurrences(self): + self.client.force_authenticate(user=self.curator) + response = self.client.get(self.url, {"project_id": self.project.pk}) + + row = next(r for r in response.data["results"] if r["name"] == "Blind set") + self.assertEqual(row["occurrences_count"], 3) + self.assertNotIn("occurrences", row) + + def test_another_project_does_not_see_it(self): + other = Project.objects.create(name="Other Project", owner=self.owner) + MLDataManager.assign_user(self.curator, other) + self.client.force_authenticate(user=self.curator) + + response = self.client.get(self.url, {"project_id": other.pk}) + + self.assertNotIn("Blind set", [row["name"] for row in response.data["results"]]) diff --git a/ami/ml/migrations/0031_algorithm_training_fields.py b/ami/ml/migrations/0031_algorithm_training_fields.py new file mode 100644 index 000000000..42d9080b9 --- /dev/null +++ b/ami/ml/migrations/0031_algorithm_training_fields.py @@ -0,0 +1,44 @@ +# Generated by Django 4.2.10 on 2026-10-08 16:04 + +import ami.ml.schemas +from django.db import migrations, models +import django_pydantic_field.fields + + +class Migration(migrations.Migration): + dependencies = [ + ("ml", "0030_detection_embedding"), + ] + + operations = [ + migrations.AddField( + model_name="algorithm", + name="trainable", + field=models.BooleanField( + default=False, + help_text="Whether the processing service reports that this algorithm can be retrained. Mirrored from the service's /info response so the UI does not have to ask it.", + ), + ), + migrations.AddField( + model_name="algorithm", + name="training_config", + field=django_pydantic_field.fields.PydanticSchemaField( + blank=True, + config=None, + default=ami.ml.schemas.AlgorithmTrainingConfig, + help_text="Settings used when retraining this algorithm.", + schema=ami.ml.schemas.AlgorithmTrainingConfig, + ), + ), + migrations.AddField( + model_name="algorithm", + name="training_info", + field=django_pydantic_field.fields.PydanticSchemaField( + blank=True, + config=None, + default=ami.ml.schemas.AlgorithmTrainingInfo, + help_text="Where this version's weights came from. Empty on a version that was never retrained.", + schema=ami.ml.schemas.AlgorithmTrainingInfo, + ), + ), + ] diff --git a/ami/ml/models/algorithm.py b/ami/ml/models/algorithm.py index 605f90861..b2f9840f6 100644 --- a/ami/ml/models/algorithm.py +++ b/ami/ml/models/algorithm.py @@ -12,8 +12,10 @@ from django.contrib.postgres.fields import ArrayField from django.db import models from django.utils.text import slugify +from django_pydantic_field import SchemaField from ami.base.models import BaseModel, BaseQuerySet +from ami.ml.schemas import AlgorithmTrainingConfig, AlgorithmTrainingInfo @typing.final @@ -249,6 +251,31 @@ class Algorithm(BaseModel): help_text=("A URI to the weights or model details. Could be a public web URL or object store path."), ) + trainable = models.BooleanField( + default=False, + help_text=( + "Whether the processing service reports that this algorithm can be retrained. " + "Mirrored from the service's /info response so the UI does not have to ask it." + ), + ) + # Seeded from the service on first registration, then owned here so an admin's edits + # are not overwritten every time /info is read. + training_config = SchemaField( + AlgorithmTrainingConfig, + default=AlgorithmTrainingConfig, + null=False, + blank=True, + help_text="Settings used when retraining this algorithm.", + ) + # Written by the training job. Every retrain creates a new version, so this records + # where one particular set of weights came from. + training_info = SchemaField( + AlgorithmTrainingInfo, + default=AlgorithmTrainingInfo, + null=False, + blank=True, + help_text="Where this version's weights came from. Empty on a version that was never retrained.", + ) category_map = models.ForeignKey( AlgorithmCategoryMap, on_delete=models.CASCADE, diff --git a/ami/ml/models/pipeline.py b/ami/ml/models/pipeline.py index b7baebd32..3ebf12e58 100644 --- a/ami/ml/models/pipeline.py +++ b/ami/ml/models/pipeline.py @@ -504,7 +504,15 @@ def get_or_create_algorithm_and_category_map( fields_to_update = { "task_type": algorithm_config.task_type, "uri": algorithm_config.uri, + # Whether the service can retrain this one. Mirrored here so the UI can offer it + # without asking the service, and so a service that drops the ability stops + # offering it too. + "trainable": algorithm_config.trainable, } + if algorithm_config.training_config and _created: + # Seeded from the service once. After that the settings are Antenna's, so an + # admin's edits are not overwritten every time the pipelines are re-registered. + fields_to_update["training_config"] = algorithm_config.training_config for field in fields_to_update: new_value = fields_to_update[field] if getattr(algo, field) != new_value: diff --git a/ami/ml/schemas.py b/ami/ml/schemas.py index 8df39267d..88ff797af 100644 --- a/ami/ml/schemas.py +++ b/ami/ml/schemas.py @@ -92,6 +92,76 @@ class AlgorithmCategoryMapResponse(pydantic.BaseModel): ) +class AlgorithmTrainingConfig(pydantic.BaseModel): + """ + How to retrain an algorithm. Declared by the processing service, editable afterwards. + + Split across two sides on purpose: Antenna reads the dataset settings when it builds + the training set, and passes the rest to the service, which owns the fitting. + """ + + # Dataset settings, used by Antenna. + min_per_species: int = pydantic.Field( + default=2, + ge=1, + description="Drop species with fewer verified crops than this. One example cannot be evaluated.", + ) + test_fraction: float = pydantic.Field( + default=0.2, + gt=0, + lt=1, + description="Share of occurrences held out for evaluation.", + ) + split_salt: str = pydantic.Field( + default="antenna-head-v1", + min_length=1, + description="Changing this reshuffles the held-out set, which makes old and new heads incomparable.", + ) + + # Fitting settings, used by the processing service. + head_type: str = pydantic.Field( + default="linear", + min_length=1, + description="Shape of the head to fit, e.g. 'linear' or 'mlp1'.", + examples=["linear", "mlp1"], + ) + epochs: int = pydantic.Field(default=300, ge=1) + learning_rate: float = pydantic.Field(default=0.01, gt=0) + weight_decay: float = pydantic.Field(default=1e-4, ge=0) + min_improvement: float = pydantic.Field( + default=0.0, + ge=0, + description="A new head must beat the current one by more than this to be worth swapping in.", + ) + + class Config: + extra = "allow" + + @classmethod + def for_run( + cls, + base: "AlgorithmTrainingConfig", + overrides: dict | None = None, + ) -> "AlgorithmTrainingConfig": + """ + The settings for one run: the algorithm's config with a job's params on top. + + Params reach this straight from the API, so they are checked here rather than read + raw at the point of use. The bounds are not housekeeping: a test_fraction of 1.5 + holds out every occurrence and leaves nothing to train on, and a negative one holds + out nothing and then scores the head on the rows it just learned. Both finish + looking like a successful run and report a number that means nothing. + + Only keys this config already carries are taken, so job settings such as + algorithm_key and taxa_list_id stay out of what is sent to the service. A setting a + service declared itself is carried on the base config, so it can be overridden too. + """ + current = base.dict() + allowed = set(cls.__fields__) | set(current) + taken = {key: value for key, value in (overrides or {}).items() if key in allowed} + return cls(**{**current, **taken}) + + class AlgorithmConfigResponse(pydantic.BaseModel): name: str key: str = pydantic.Field( @@ -115,6 +185,14 @@ class AlgorithmConfigResponse(pydantic.BaseModel): default=None, description="A URI to the weight or model details, could be a public web URL or object store path.", ) + trainable: bool = pydantic.Field( + default=False, + description="Whether this service can retrain this algorithm from a dataset of verified crops.", + ) + training_config: AlgorithmTrainingConfig | None = pydantic.Field( + default=None, + description="The settings this service would retrain with. Seeds the algorithm's config in Antenna.", + ) category_map: AlgorithmCategoryMapResponse | None = None class Config: @@ -393,3 +471,151 @@ class AsyncPipelineRegistrationRequest(pydantic.BaseModel): processing_service_name: str 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. + + Every retrain produces a new algorithm version, so this is the record of where a + particular set of weights came from. + """ + + trained_at: datetime.datetime | None = None + dataset_url: str | None = pydantic.Field(default=None, description="The exact training set this was fitted on.") + occurrence_set_id: int | None = None + occurrence_set_name: str | None = pydantic.Field( + default=None, + description="Named as well as numbered, so a version still says what it learned from if the set is deleted.", + ) + dataset_rows: int | None = None + dataset_classes: int | None = None + metrics: dict = pydantic.Field(default_factory=dict, description="Scores on the held-out split.") + previous_metrics: dict = pydantic.Field( + default_factory=dict, description="What the version it was compared against scored on the same rows." + ) + parent_algorithm_key: str | None = pydantic.Field( + default=None, description="The version this one was trained to beat." + ) + job_id: int | None = None + warnings: list[str] = pydantic.Field(default_factory=list) + + class Config: + extra = "allow" + + +class NamedReference(pydantic.BaseModel): + """ + A row a training set came from. + + Carries the name as well as the id so a dataset still says where it came from after the + row it points at is deleted. + """ + + id: int + name: str + + +class TrainedAlgorithmReference(AlgorithmReference): + """The algorithm a training set was built for, at the version it was then.""" + + version: int | None = None + + +class TrainingDatasetSettings(pydantic.BaseModel): + """The settings a training set was built under, so a later run can be made to match.""" + + min_per_species: int + split_salt: str + test_fraction: float + split_grouped_by: str = pydantic.Field( + default="occurrence", + description=( + "What the train/test split is grouped by. Splitting per detection puts " + "near-identical crops on both sides." + ), + ) + + +class TrainingDatasetMetadata(pydantic.BaseModel): + """ + What a training set is, written inside the file and echoed back by the service. + + Declared rather than assembled as a dict so the shape is checked once, here, instead of + at each reader. + """ + + url: str + project: NamedReference + algorithm: TrainedAlgorithmReference + dimensions: int = pydantic.Field(description="Width of one feature vector, read from the stored vectors.") + dtype: str + classes: list[str] = pydantic.Field(description="The head's output classes, in order.") + counts: dict[str, int] = pydantic.Field(default_factory=dict, description="Verified crops per class.") + taxa_list: NamedReference | None = None + occurrence_set: NamedReference | None = None + classes_without_verified_data: list[str] = pydantic.Field(default_factory=list) + dropped_species: list[str] = pydantic.Field(default_factory=list) + rows: int + train: int + test: int + verified_detections_without_embedding: int = 0 + settings: TrainingDatasetSettings diff --git a/ami/ml/serializers.py b/ami/ml/serializers.py index e7e9e6aaf..0ece73409 100644 --- a/ami/ml/serializers.py +++ b/ami/ml/serializers.py @@ -7,7 +7,7 @@ from .models.pipeline import Pipeline, PipelineStage from .models.processing_service import ProcessingService from .models.project_pipeline_config import ProjectPipelineConfig -from .schemas import PipelineConfigResponse +from .schemas import AlgorithmTrainingConfig, AlgorithmTrainingInfo, PipelineConfigResponse class AlgorithmCategoryMapSerializer(DefaultSerializer): @@ -30,6 +30,10 @@ class Meta: class AlgorithmSerializer(DefaultSerializer): category_map = MinimalCategoryMapNestedSerializer(read_only=True, source="category_map_id") + # Pydantic-backed model fields have to be declared, the way the job serializers do. + # A bare name in Meta.fields raises ImproperlyConfigured at import time. + training_config = SchemaField(schema=AlgorithmTrainingConfig, required=False) + training_info = SchemaField(schema=AlgorithmTrainingInfo, read_only=True) class Meta: model = Algorithm @@ -43,6 +47,9 @@ class Meta: "version", "version_name", "task_type", + "trainable", + "training_config", + "training_info", "category_map", "category_count", "created_at", @@ -168,3 +175,22 @@ def get_projects(self, obj): class PipelineRegistrationSerializer(serializers.Serializer): processing_service_name = serializers.CharField() pipelines = SchemaField(schema=list[PipelineConfigResponse], default=[]) + + +class TrainingDataRowSerializer(serializers.Serializer): + """ + One verified crop: its human label and the embedding to train on. + + Not a ModelSerializer — the useful fields live on the related occurrence and on the + pgvector column, and the split is computed, so building the dict directly is clearer. + """ + + def to_representation(self, instance): + from ami.ml import training + + return training.row_as_dict( + instance, + salt=self.context.get("split_salt", training.DEFAULT_SPLIT_SALT), + test_fraction=self.context.get("test_fraction", training.DEFAULT_TEST_FRACTION), + include_features=self.context.get("include_features", True), + ) diff --git a/ami/ml/test_training.py b/ami/ml/test_training.py new file mode 100644 index 000000000..b75eede40 --- /dev/null +++ b/ami/ml/test_training.py @@ -0,0 +1,518 @@ +""" +Building a training set from an occurrence set. + +A run learns from a named set rather than from whatever happened to be verified that day, +so these tests are mostly about the set deciding what goes in. +""" + +import datetime +import uuid + +from django.test import TestCase +from django.urls import reverse +from rest_framework.test import APIClient + +from ami.main.models import ( + Classification, + Deployment, + Detection, + Event, + Identification, + Occurrence, + OccurrenceSet, + Project, + S3StorageSource, + SourceImage, + Taxon, + TaxonRank, +) +from ami.ml.models import Algorithm, AlgorithmCategoryMap +from ami.ml.models.embedding import DetectionEmbedding +from ami.ml.schemas import NamedReference +from ami.ml.training import NotEnoughVerifiedData, build_training_dataset, verified_training_rows +from ami.users.models import User +from ami.users.roles import ProjectManager + +VECTOR_LENGTH = 8 + + +class TrainingSetFixture(TestCase): + """ + Occurrences, identifications and vectors built directly. + + The capture fixtures time their images from datetime.now() and group events by reading + image dimensions from storage, neither of which is deterministic in a test. + """ + + def setUp(self) -> None: + self.user = User.objects.create_user(email="curator@example.com", password="testpass123") + self.project = Project.objects.create(name="Training Project", owner=self.user) + storage = S3StorageSource.objects.create(name="src", project=self.project, bucket="test") + self.deployment = Deployment.objects.create(name="station", project=self.project, data_source=storage) + night = datetime.datetime(2026, 6, 1, 22, 0, 0, tzinfo=datetime.timezone.utc) + self.event = Event.objects.create(deployment=self.deployment, project=self.project, start=night) + self.night = night + + self.algorithm = Algorithm.objects.create(name="Backbone", key="backbone", trainable=True) + self.taxa = [ + Taxon.objects.create(name=f"Trainus {name}", rank=TaxonRank.SPECIES.name) for name in ("alpha", "beta") + ] + for taxon in self.taxa: + taxon.projects.add(self.project) + self.algorithm.category_map = AlgorithmCategoryMap.objects.create( + labels=[t.name for t in self.taxa], + data=[{"index": i, "label": t.name} for i, t in enumerate(self.taxa)], + version="training-test", + ) + self.algorithm.save() + + def make_occurrence(self, taxon: Taxon, with_vector: bool = True) -> Occurrence: + """One verified occurrence, optionally carrying a feature vector.""" + image = SourceImage.objects.create( + deployment=self.deployment, + project=self.project, + event=self.event, + timestamp=self.night, + path=f"test/{uuid.uuid4().hex[:8]}.jpg", + ) + detection = Detection.objects.create(source_image=image, timestamp=self.night, bbox=[0.1, 0.1, 0.2, 0.2]) + Classification.objects.create(detection=detection, taxon=taxon, score=0.9, timestamp=self.night) + occurrence = detection.associate_new_occurrence() + Identification.objects.create(occurrence=occurrence, taxon=taxon, user=self.user) + if with_vector: + DetectionEmbedding.objects.create( + detection=detection, + algorithm=self.algorithm, + project=self.project, + vector=[0.1] * VECTOR_LENGTH, + ) + occurrence.refresh_from_db() + return occurrence + + def make_set(self, name: str, occurrences: list[Occurrence]) -> OccurrenceSet: + occurrence_set = OccurrenceSet.objects.create(name=name) + occurrence_set.projects.add(self.project) + occurrence_set.occurrences.set(occurrences) + return occurrence_set + + +class TestTheSetDecidesWhatIsLearnedFrom(TrainingSetFixture): + def test_only_the_sets_occurrences_are_used(self): + inside = [self.make_occurrence(self.taxa[0]), self.make_occurrence(self.taxa[1])] + outside = self.make_occurrence(self.taxa[0]) + occurrence_set = self.make_set("Spring review", inside) + + rows = verified_training_rows(self.project, self.algorithm, occurrence_set) + used = {row.detection.occurrence_id for row in rows} + + self.assertEqual(used, {o.pk for o in inside}) + self.assertNotIn(outside.pk, used) + + def test_without_a_set_every_verified_occurrence_is_used(self): + """The helper still answers without one; it is the job that insists on a set.""" + occurrences = [self.make_occurrence(self.taxa[0]), self.make_occurrence(self.taxa[1])] + + rows = verified_training_rows(self.project, self.algorithm) + + self.assertEqual({row.detection.occurrence_id for row in rows}, {o.pk for o in occurrences}) + + def test_the_dataset_says_which_set_it_came_from(self): + occurrences = [self.make_occurrence(self.taxa[0]) for _ in range(4)] + occurrences += [self.make_occurrence(self.taxa[1]) for _ in range(4)] + occurrence_set = self.make_set("Spring review", occurrences) + + # Half held out: the split is a hash of the occurrence, so a handful of rows at the + # usual fifth can land entirely on one side, which the build refuses. + dataset = build_training_dataset( + project=self.project, + algorithm=self.algorithm, + min_per_species=1, + test_fraction=0.5, + occurrence_set=occurrence_set, + ) + + self.assertEqual( + dataset["metadata"].occurrence_set, + NamedReference(id=occurrence_set.pk, name="Spring review"), + ) + self.assertEqual(dataset["metadata"].rows, 8) + + def test_an_empty_set_has_nothing_to_train_on(self): + self.make_occurrence(self.taxa[0]) + empty = self.make_set("Nothing in here", []) + + with self.assertRaises(NotEnoughVerifiedData): + build_training_dataset( + project=self.project, algorithm=self.algorithm, min_per_species=1, occurrence_set=empty + ) + + def test_the_vector_width_comes_from_the_stored_vectors(self): + """Nothing assumes a width: a backbone of another size is read, not truncated.""" + occurrences = [self.make_occurrence(self.taxa[0]) for _ in range(6)] + occurrence_set = self.make_set("Spring review", occurrences) + + dataset = build_training_dataset( + project=self.project, + algorithm=self.algorithm, + min_per_species=1, + test_fraction=0.5, + occurrence_set=occurrence_set, + ) + + self.assertEqual(dataset["metadata"].dimensions, VECTOR_LENGTH) + + +class TestTheSummaryTheFormShows(TrainingSetFixture): + """ + The numbers the training job form puts in front of someone before they start a run. + + A form that asked for a split ratio without saying how many occurrences it would split + would be asking blind, so these cover what changes when the set or the ratio changes. + """ + + def setUp(self) -> None: + super().setUp() + ProjectManager.assign_user(self.user, self.project) + self.client = APIClient() + self.client.force_authenticate(user=self.user) + self.url = reverse("api:training-data-summary") + + def _summary(self, **params): + response = self.client.get( + self.url, + {"project_id": self.project.pk, "algorithm": self.algorithm.key, **params}, + ) + self.assertEqual(response.status_code, 200, response.data) + return response.data + + def test_without_a_set_it_counts_every_verified_occurrence(self): + for taxon in self.taxa: + self.make_occurrence(taxon) + + summary = self._summary() + + self.assertIsNone(summary["occurrence_set"]) + self.assertEqual(summary["occurrences"], 2) + self.assertEqual(summary["rows"], 2) + + def test_a_set_narrows_the_counts_and_is_named_back(self): + mine = self.make_set("Mine", [self.make_occurrence(self.taxa[0])]) + self.make_occurrence(self.taxa[1]) + + summary = self._summary(occurrence_set=mine.pk) + + self.assertEqual(summary["occurrence_set"], {"id": mine.pk, "name": "Mine"}) + self.assertEqual(summary["occurrences"], 1) + + def test_a_set_from_another_project_is_not_counted(self): + other = Project.objects.create(name="Someone else", owner=self.user) + theirs = OccurrenceSet.objects.create(name="Theirs") + theirs.projects.add(other) + + response = self.client.get( + self.url, + {"project_id": self.project.pk, "algorithm": self.algorithm.key, "occurrence_set": theirs.pk}, + ) + + 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) + + mostly_train = self._summary(test_fraction=0.01) + mostly_test = self._summary(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( + self.url, + {"project_id": self.project.pk, "algorithm": self.algorithm.key, "test_fraction": 1.5}, + ) + + self.assertEqual(response.status_code, 400) + + def test_species_without_enough_crops_are_reported_as_dropped(self): + self.make_occurrence(self.taxa[0]) + self.make_occurrence(self.taxa[1]) + self.make_occurrence(self.taxa[1]) + + summary = self._summary(min_per_species=2) + + self.assertEqual(summary["classes"], 2) + self.assertEqual(summary["trainable_classes"], 1) + self.assertEqual(summary["dropped_species"], [self.taxa[0].name]) + + def test_someone_who_cannot_retrain_cannot_read_the_training_set(self): + """ + A non-draft project is readable by anyone, so visibility alone would hand every + verified label and vector to any account. The permission is the gate. + """ + self.make_occurrence(self.taxa[0]) + outsider = User.objects.create_user(email="outsider@example.com", password="testpass123") + self.client.force_authenticate(user=outsider) + + response = self.client.get(self.url, {"project_id": self.project.pk, "algorithm": self.algorithm.key}) + + self.assertEqual(response.status_code, 403) + + +class TestTheJobsOccurrenceSet(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_learns_from_everything_verified(self): + from ami.jobs.models import TrainClassifierJob + + self.assertIsNone(TrainClassifierJob.target_occurrence_set(self._job())) + + 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) + + +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) + + +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/__init__.py b/ami/ml/training/__init__.py new file mode 100644 index 000000000..91bab508e --- /dev/null +++ b/ami/ml/training/__init__.py @@ -0,0 +1,69 @@ +""" +Retraining a classifier head from the species people have verified. + +``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 ( + DEFAULT_SPLIT_SALT, + DEFAULT_TEST_FRACTION, + SPLIT_TEST, + SPLIT_TRAIN, + SPLITS, + NotEnoughVerifiedData, + build_training_dataset, + count_missing_embeddings, + label_counts, + row_as_dict, + species_with_enough_examples, + split_for, + 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, + progress_url_for, + 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", + "progress_url_for", + "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/dataset.py b/ami/ml/training/dataset.py new file mode 100644 index 000000000..d81a1e8dd --- /dev/null +++ b/ami/ml/training/dataset.py @@ -0,0 +1,341 @@ +""" +Build a classifier-head training set out of what people verified, and write it to storage. + +Two halves of one job: deciding what counts as training data — which occurrences a person +identified, which species have enough of them, which rows are held out — and turning that +into a single file. + +It is a file because a service is handed a URL, not 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. +""" + +import hashlib +import logging +import pathlib +import tempfile +import typing +from typing import TYPE_CHECKING + +import numpy as np +from django.conf import settings +from django.core.files.storage import default_storage +from django.db.models import QuerySet +from django.utils.text import slugify + +from ami.main.models import DEFAULT_EMBEDDING_KEY, Identification, Project +from ami.ml.models.algorithm import Algorithm +from ami.ml.models.embedding import DetectionEmbedding +from ami.ml.schemas import NamedReference, TrainedAlgorithmReference, TrainingDatasetMetadata, TrainingDatasetSettings + +if TYPE_CHECKING: + from ami.jobs.models import Job + from ami.main.models import OccurrenceSet, TaxaList + +logger = logging.getLogger(__name__) + + +DEFAULT_SPLIT_SALT = "antenna-head-v1" +DEFAULT_TEST_FRACTION = 0.2 + +# The two sides of the split. Named so the API, the dataset builder and the export +# cannot drift apart on spelling. +SPLIT_TRAIN = "train" +SPLIT_TEST = "test" +SPLITS = (SPLIT_TRAIN, SPLIT_TEST) + + +def split_for( + occurrence_id: int, + salt: str = DEFAULT_SPLIT_SALT, + test_fraction: float = DEFAULT_TEST_FRACTION, +) -> str: + """ + Assign an occurrence to "train" or "test", the same way every time. + + Grouped by occurrence, not detection: an occurrence is one insect across several + frames, so a detection-level split puts near-identical crops on both sides and + overstates accuracy. Hash-based, not random, so the test set does not move when new + data arrives — an eval set that drifts cannot compare two heads. + """ + digest = hashlib.sha256(f"{salt}:{occurrence_id}".encode()).hexdigest() + return SPLIT_TEST if (int(digest[:8], 16) / 0xFFFFFFFF) < test_fraction else SPLIT_TRAIN + + +def verified_occurrence_ids(project: Project, occurrence_set: "OccurrenceSet | None" = None) -> QuerySet: + """ + Occurrences a person has identified and not taken back. + + Narrowed to an occurrence set when one is given. A set is fixed once created, so a run + that names one can be repeated and says for itself what it learned from; without one the + pool is whatever happened to be verified that day, which no later reader can reconstruct. + """ + identifications = Identification.objects.filter( + withdrawn=False, + occurrence__project=project, + occurrence__determination__isnull=False, + ) + if occurrence_set is not None: + identifications = identifications.filter(occurrence__evaluation_sets=occurrence_set) + return identifications.order_by().values_list("occurrence_id", flat=True).distinct() + + +def verified_training_rows( + project: Project, algorithm: Algorithm, occurrence_set: "OccurrenceSet | None" = None +) -> QuerySet[DetectionEmbedding]: + """ + Embeddings whose detection sits under a verified occurrence. + + Constrained to one algorithm on purpose: vectors from different backbones are in + different spaces and must never be mixed into one training set. + """ + return ( + DetectionEmbedding.objects.filter( + algorithm=algorithm, + key=DEFAULT_EMBEDDING_KEY, + detection__occurrence_id__in=verified_occurrence_ids(project, occurrence_set), + detection__occurrence__determination__isnull=False, + ) + .select_related("detection__occurrence__determination") + .order_by("pk") + ) + + +def count_missing_embeddings( + project: Project, algorithm: Algorithm, occurrence_set: "OccurrenceSet | None" = None +) -> int: + """Verified detections this algorithm has never embedded. They need a pipeline re-run.""" + from ami.main.models import Detection + + return ( + Detection.objects.filter(occurrence_id__in=verified_occurrence_ids(project)) + .exclude(embeddings__algorithm=algorithm) + .count() + ) + + +def label_counts( + project: Project, algorithm: Algorithm, occurrence_set: "OccurrenceSet | None" = None +) -> dict[str, int]: + """Verified crops per species, for the chosen algorithm.""" + from django.db.models import Count + + rows = ( + verified_training_rows(project, algorithm, occurrence_set) + .order_by() + .values("detection__occurrence__determination__name") + .annotate(n=Count("pk")) + ) + return {r["detection__occurrence__determination__name"]: r["n"] for r in rows} + + +def species_with_enough_examples(counts: dict[str, int], minimum: int) -> set[str]: + """A class with a single example cannot be both trained on and evaluated.""" + return {name for name, n in counts.items() if name and n >= minimum} + + +def row_as_dict( + embedding: DetectionEmbedding, + salt: str = DEFAULT_SPLIT_SALT, + test_fraction: float = DEFAULT_TEST_FRACTION, + include_features: bool = True, +) -> dict[str, typing.Any]: + """One training row, in the shape the training-data API returns.""" + occurrence = embedding.detection.occurrence + row = { + "detection_id": embedding.detection_id, + "occurrence_id": occurrence.pk if occurrence else None, + "label": occurrence.determination.name if occurrence and occurrence.determination else None, + "label_id": occurrence.determination_id if occurrence else None, + "split": split_for(occurrence.pk, salt, test_fraction) if occurrence else None, + } + if include_features: + row["features"] = embedding.vector.to_list() + return row + + +# Vectors are written as float16 because that is exactly what Postgres stores (halfvec), +# so nothing is lost. Measured on 10k rows: npz float16 is 21 MB against 227 MB of JSON, +# and 1.1s to build against 29s. +DATASET_DTYPE = np.float16 + +DATASET_DIRECTORY = "training" + +# Rows are pulled from the database in batches so a large project does not have to fit +# every vector in memory at once. +FETCH_BATCH_SIZE = 2000 + + +class NotEnoughVerifiedData(Exception): + """The project does not hold enough verified labels to train and evaluate on.""" + + +def build_training_dataset( + project: Project, + algorithm: Algorithm, + min_per_species: int = 2, + split_salt: str = DEFAULT_SPLIT_SALT, + test_fraction: float = DEFAULT_TEST_FRACTION, + taxa_list: "TaxaList | None" = None, + occurrence_set: "OccurrenceSet | None" = None, + job: "Job | None" = None, +) -> dict[str, typing.Any]: + """ + Collect verified labels and their embeddings, and save them as one npz file. + + Returns the storage path, the URL a service can fetch, and the metadata describing + what went in. Raises NotEnoughVerifiedData rather than writing a file nothing can be + trained on. + """ + counts = label_counts(project, algorithm, occurrence_set) + taxa_list = taxa_list or project.default_taxa_list + + if taxa_list: + # The taxa list decides what the head can predict; the verified crops only decide + # how well it predicts each one. Without this the head shrinks to whatever happened + # to be verified, which silently narrows the pipeline. + classes = sorted(taxa_list.taxa.exclude(name="").values_list("name", flat=True)) + if not classes: + raise NotEnoughVerifiedData(f"Taxa list '{taxa_list}' is empty, so there is nothing to train.") + keep = set(classes) + outside = sorted(name for name in counts if name and name not in keep) + else: + keep = species_with_enough_examples(counts, min_per_species) + if not keep: + raise NotEnoughVerifiedData( + f"No species in '{project.name}' has at least {min_per_species} verified crops with an " + f"embedding from {algorithm.key}. Verify more occurrences, or re-run the pipeline so the " + "verified detections get embeddings." + ) + classes = sorted(keep) + outside = sorted(set(counts) - keep) + + class_index = {name: i for i, name in enumerate(classes)} + + rows = verified_training_rows(project, algorithm, occurrence_set) + total = rows.count() + + # Asked rather than assumed: a vector's width belongs to the model that produced it, and + # the column is unsized so two algorithms may differ. + dimensions = DetectionEmbedding.objects.stored_length(algorithm.pk) + if dimensions is None: + raise NotEnoughVerifiedData( + f"No detection has a feature vector from '{algorithm.key}' yet, so there is nothing to train on." + ) + features = np.zeros((total, dimensions), dtype=DATASET_DTYPE) + labels = np.zeros(total, dtype=np.int64) + detection_ids = np.zeros(total, dtype=np.int64) + occurrence_ids = np.zeros(total, dtype=np.int64) + splits: list[str] = [] + + kept = 0 + for embedding in rows.iterator(chunk_size=FETCH_BATCH_SIZE): + occurrence = embedding.detection.occurrence + if not occurrence or not occurrence.determination: + continue + name = occurrence.determination.name + if name not in class_index: + continue + features[kept] = np.asarray(embedding.vector.to_list(), dtype=DATASET_DTYPE) + labels[kept] = class_index[name] + detection_ids[kept] = embedding.detection_id + occurrence_ids[kept] = occurrence.pk + splits.append(split_for(occurrence.pk, split_salt, test_fraction)) + kept += 1 + + if not kept: + raise NotEnoughVerifiedData("No verified detection has an embedding from this algorithm yet.") + + features = features[:kept] + labels = labels[:kept] + detection_ids = detection_ids[:kept] + occurrence_ids = occurrence_ids[:kept] + split_array = np.array(splits) + + n_train = int((split_array == "train").sum()) + n_test = int((split_array == "test").sum()) + if not n_train or not n_test: + raise NotEnoughVerifiedData( + f"The split left {n_train} training and {n_test} held-out rows. Both sides need rows " + "before a new head can be compared against the current one." + ) + + file_path = dataset_path(project, algorithm, job.pk if job else None) + # Declared, not assembled: the shape is checked here rather than at each reader, and + # the service echoes it back in its result. + metadata = TrainingDatasetMetadata( + url=f"{settings.MEDIA_URL}{file_path}", + project=NamedReference(id=project.pk, name=project.name), + algorithm=TrainedAlgorithmReference(key=algorithm.key, name=algorithm.name, version=algorithm.version), + dimensions=dimensions, + dtype=np.dtype(DATASET_DTYPE).name, + classes=classes, + # .get(): with a taxa list, a class can legitimately have no verified crops yet. + counts={name: counts.get(name, 0) for name in classes}, + taxa_list=NamedReference(id=taxa_list.pk, name=taxa_list.name) if taxa_list else None, + occurrence_set=(NamedReference(id=occurrence_set.pk, name=occurrence_set.name) if occurrence_set else None), + classes_without_verified_data=sorted(name for name in classes if not counts.get(name)), + dropped_species=outside, + rows=kept, + train=n_train, + test=n_test, + verified_detections_without_embedding=count_missing_embeddings(project, algorithm, occurrence_set), + settings=TrainingDatasetSettings( + min_per_species=min_per_species, + split_salt=split_salt, + test_fraction=test_fraction, + ), + ) + + file_path = _save( + file_path=file_path, + arrays={ + "features": features, + "labels": labels, + "detection_ids": detection_ids, + "occurrence_ids": occurrence_ids, + "split": split_array, + "classes": np.array(classes), + }, + metadata=metadata, + ) + + file_url = metadata.url + logger.info(f"Wrote training dataset with {kept} rows over {len(classes)} species to {file_path}") + return {"path": file_path, "url": file_url, "metadata": metadata} + + +def dataset_path(project: Project, algorithm: Algorithm, job_id: int | None) -> str: + """Where this project's training set for this algorithm is written.""" + stem = f"{slugify(project.name)}-{slugify(algorithm.key)}" + suffix = f"job-{job_id}" if job_id else "manual" + return f"{DATASET_DIRECTORY}/{stem}-{suffix}.npz" + + +def _save( + file_path: str, + arrays: dict[str, np.ndarray], + metadata: TrainingDatasetMetadata, +) -> str: + """ + Write the npz to storage and return its path. + + default_storage is the local filesystem in development and the project's S3 bucket in + production, so this follows wherever captures already live without a special case. + """ + with tempfile.TemporaryDirectory() as tmp: + local = pathlib.Path(tmp) / "dataset.npz" + # Uncompressed: embeddings are close to random, so compression buys about 10 per + # cent for real CPU. Metadata rides inside the archive so the file is self-describing. + np.savez(local, metadata=np.array(metadata.json()), **arrays) + if default_storage.exists(file_path): + # A re-run of the same job replaces its dataset instead of piling up copies. + default_storage.delete(file_path) + with open(local, "rb") as f: + saved_path = default_storage.save(file_path, f) + + if saved_path != file_path: + # The metadata inside the archive records the URL it was written to, so a rename + # by the storage backend would leave the file describing somewhere it is not. + logger.warning(f"Storage wrote the training set to {saved_path}, not {file_path}.") + return saved_path diff --git a/ami/ml/training/service.py b/ami/ml/training/service.py new file mode 100644 index 000000000..f684708f6 --- /dev/null +++ b/ami/ml/training/service.py @@ -0,0 +1,200 @@ +""" +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, TrainingRequest +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 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. + + 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) -> None: + """ + Ask a processing service to retrain a head. + + 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 {}) + 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. + 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=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}") + + +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 diff --git a/ami/ml/views.py b/ami/ml/views.py index 63e460af6..516aa3444 100644 --- a/ami/ml/views.py +++ b/ami/ml/views.py @@ -1,24 +1,31 @@ import logging +import pydantic from django.db import transaction from django.db.models import Prefetch from django.db.models.query import QuerySet +from django.http import Http404 from django.utils.text import slugify from drf_spectacular.utils import extend_schema from rest_framework import exceptions as api_exceptions -from rest_framework import mixins, status, viewsets +from rest_framework import mixins, serializers, status, viewsets from rest_framework.decorators import action +from rest_framework.permissions import IsAuthenticated from rest_framework.request import Request from rest_framework.response import Response +from ami.base.pagination import TrainingDataPagination from ami.base.permissions import ProjectPipelineConfigPermission +from ami.base.serializers import SingleParamSerializer from ami.base.views import ProjectMixin from ami.main.api.schemas import project_id_doc_param from ami.main.api.views import DefaultViewSet -from ami.main.models import Project, SourceImage -from ami.ml.schemas import PipelineRegistrationResponse +from ami.main.models import OccurrenceSet, Project, SourceImage +from ami.ml import training +from ami.ml.schemas import AlgorithmTrainingConfig, PipelineRegistrationResponse from .models.algorithm import Algorithm, AlgorithmCategoryMap +from .models.embedding import DetectionEmbedding from .models.pipeline import Pipeline from .models.processing_service import ProcessingService from .models.project_pipeline_config import ProjectPipelineConfig @@ -28,6 +35,7 @@ PipelineRegistrationSerializer, PipelineSerializer, ProcessingServiceSerializer, + TrainingDataRowSerializer, ) logger = logging.getLogger(__name__) @@ -40,7 +48,7 @@ class AlgorithmViewSet(DefaultViewSet, ProjectMixin): queryset = Algorithm.objects.all() serializer_class = AlgorithmSerializer - filterset_fields = ["name", "version"] + filterset_fields = ["name", "version", "trainable"] ordering_fields = [ "id", "created_at", @@ -298,3 +306,216 @@ def create(self, request, *args, **kwargs): processing_service.mark_seen(live=True) return Response(response.dict(), status=status.HTTP_201_CREATED) + + +class TrainingDataViewSet(ProjectMixin, mixins.ListModelMixin, viewsets.GenericViewSet): + """ + Verified detections and their embeddings, for retraining a classifier head. + + A head is trained on embeddings, not pixels, and the backbone that produced them is + frozen. So a trainer can pull this and fit a new head without touching the images. + + Requires `project_id` and `algorithm` (an algorithm key). Constraining to one + algorithm is not optional: vectors from different backbones are in different spaces. + + `occurrence_set` narrows this to one saved set. Without it the rows are every verified + occurrence in the project, which is what a retrain does by default. + + GET /api/v2/ml/training-data/?project_id=3&algorithm= + GET /api/v2/ml/training-data/summary/?project_id=3&algorithm=&occurrence_set=7 + """ + + queryset = DetectionEmbedding.objects.none() + serializer_class = TrainingDataRowSerializer + require_project = True + # The project permission is checked per request, in _get_project_for_training(). + # ObjectPermission's has_object_permission only runs for detail actions, and both + # actions here are detail=False, so a viewset-level class would never see the project. + permission_classes = [IsAuthenticated] + filter_backends: list = [] + pagination_class = TrainingDataPagination + + def _get_project_for_training(self) -> Project: + """ + The requested project, if this user may retrain its classifier. + + Project visibility is not enough of a gate here. A non-draft project is readable by + anyone, and what this returns is every verified label in it together with the vector + for each crop: the whole training set, not a page of occurrences. It is gated on the + same permission the training job is, so the people who can start a run are the ones + who can see what it would learn from. + """ + project = self.get_active_project() + if not project: + raise Http404("Project not found.") + if not self.request.user.has_perm(Project.Permissions.RUN_TRAIN_CLASSIFIER_JOB, project): + raise api_exceptions.PermissionDenied("You do not have permission to retrain classifiers in this project.") + return project + + def _get_algorithm(self) -> Algorithm: + # Cached because the split settings fall back to this algorithm's own training + # config, so one request asks for it from more than one place. + if getattr(self, "_algorithm", None) is None: + key = SingleParamSerializer[str].clean( + "algorithm", + serializers.CharField( + required=True, + help_text="Key of the algorithm whose embeddings to train on.", + ), + self.request.query_params, + ) + algorithm = Algorithm.objects.filter(key=key).first() + if not algorithm: + raise api_exceptions.NotFound(f"No algorithm with key '{key}'.") + self._algorithm = algorithm + return self._algorithm + + def _get_occurrence_set(self, project: Project) -> OccurrenceSet | None: + """ + The set to count, or None for every verified occurrence in the project. + + Optional, because retraining on everything verified so far is the common case and + the form shows those numbers before anyone has chosen a set. + """ + occurrence_set_id = SingleParamSerializer[int].clean( + "occurrence_set", + serializers.IntegerField(required=False, allow_null=True, default=None), + self.request.query_params, + ) + if not occurrence_set_id: + return None + # Scoped to the project: a set id from elsewhere must not report this project's + # verified labels, or any other project's. + occurrence_set = OccurrenceSet.objects.for_project(project).filter(pk=occurrence_set_id).first() + if not occurrence_set: + raise api_exceptions.NotFound(f"No occurrence set with id {occurrence_set_id} in this project.") + return occurrence_set + + def _get_split_settings(self) -> AlgorithmTrainingConfig: + """ + The settings these numbers are computed under. + + Defaulted to the algorithm's own training config rather than to the module defaults, + so a form that shows the stats before anything is edited shows what an actual run + would do. Anything the caller passes wins, which is how the form previews a changed + split ratio without starting a job. + """ + config = self._get_algorithm().training_config + params = self.request.query_params + overrides = { + "split_salt": SingleParamSerializer[str].clean( + "split_salt", + serializers.CharField(required=False, allow_null=True, default=None), + params, + ), + "test_fraction": SingleParamSerializer[float].clean( + "test_fraction", + serializers.FloatField(required=False, allow_null=True, default=None), + params, + ), + "min_per_species": SingleParamSerializer[int].clean( + "min_per_species", + serializers.IntegerField(required=False, allow_null=True, default=None), + params, + ), + } + # A parameter that was not sent must not override the config with None. + overrides = {name: value for name, value in overrides.items() if value is not None} + try: + return AlgorithmTrainingConfig.for_run(config, overrides) + except pydantic.ValidationError as e: + # The same bounds a job is held to, refused here rather than reported as stats + # computed under settings no run could use. + raise api_exceptions.ValidationError({"detail": e.errors()}) + + def get_queryset(self) -> QuerySet[DetectionEmbedding]: + project = self._get_project_for_training() + qs = training.verified_training_rows(project, self._get_algorithm(), self._get_occurrence_set(project)) + + self._split_filter = SingleParamSerializer[str].clean( + "split", + serializers.ChoiceField(choices=list(training.SPLITS), required=False, allow_null=True, default=None), + self.request.query_params, + ) + return qs + + def get_serializer_context(self): + context = super().get_serializer_context() + config = self._get_split_settings() + context["split_salt"] = config.split_salt + context["test_fraction"] = config.test_fraction + # Not url_boolean_param: it returns `value or default`, so a default of True can + # never be turned off. + context["include_features"] = SingleParamSerializer[bool].clean( + "include_features", + serializers.BooleanField(required=False, default=True), + self.request.query_params, + ) + return context + + @extend_schema(parameters=[project_id_doc_param]) + def list(self, request, *args, **kwargs): + response = super().list(request, *args, **kwargs) + split = getattr(self, "_split_filter", None) + if split: + # Filtering by split in Python rather than SQL: the assignment is a hash of the + # occurrence id, which Postgres cannot compute. Callers that need whole splits + # should page through everything and group client-side. + results = [row for row in response.data["results"] if row["split"] == split] + response.data["results"] = results + return response + + @extend_schema(parameters=[project_id_doc_param]) + @action(detail=False, methods=["get"]) + def summary(self, request, *args, **kwargs): + """ + Counts only, for what a retrain would learn from right now. + + This is what the training job form shows before anyone starts a run: how much + verified data there is, how many species of it are usable, and how the split would + fall. Cheap enough to re-fetch whenever the chosen set or the split ratio changes. + """ + project = self._get_project_for_training() + algorithm = self._get_algorithm() + occurrence_set = self._get_occurrence_set(project) + config = self._get_split_settings() + + counts = training.label_counts(project, algorithm, occurrence_set) + kept = training.species_with_enough_examples(counts, config.min_per_species) + rows = training.verified_training_rows(project, algorithm, occurrence_set) + + splits = {name: 0 for name in training.SPLITS} + occurrences = set() + for occurrence_id in rows.values_list("detection__occurrence_id", flat=True): + splits[training.split_for(occurrence_id, config.split_salt, config.test_fraction)] += 1 + occurrences.add(occurrence_id) + + return Response( + { + "project": {"id": project.pk, "name": project.name}, + "algorithm": {"key": algorithm.key, "name": algorithm.name, "version": algorithm.version}, + # None means every verified occurrence in the project, which is what a run + # without a set does. + "occurrence_set": ({"id": occurrence_set.pk, "name": occurrence_set.name} if occurrence_set else None), + "dimensions": DetectionEmbedding.objects.stored_length(algorithm.pk), + "rows": sum(counts.values()), + "occurrences": len(occurrences), + "classes": len(counts), + # What the head would actually learn, once species with too few crops are + # dropped. Showing only the raw class count overstates the run. + "trainable_classes": len(kept), + "dropped_species": sorted(name for name in counts if name and name not in kept), + "counts": dict(sorted(counts.items(), key=lambda kv: -kv[1])), + "train": splits[training.SPLIT_TRAIN], + "test": splits[training.SPLIT_TEST], + "verified_detections_without_embedding": training.count_missing_embeddings( + project, algorithm, occurrence_set + ), + "settings": { + "min_per_species": config.min_per_species, + "split_salt": config.split_salt, + "test_fraction": config.test_fraction, + "split_grouped_by": "occurrence", + }, + } + ) diff --git a/ami/users/roles.py b/ami/users/roles.py index 718146e41..d05aa209a 100644 --- a/ami/users/roles.py +++ b/ami/users/roles.py @@ -146,6 +146,9 @@ class MLDataManager(Role): Project.Permissions.SYNC_DEPLOYMENT, Project.Permissions.RUN_REGROUP_EVENTS_JOB, Project.Permissions.RUN_DATA_EXPORT_JOB, + # Retraining a head is ML data work, so it sits with the other run_*_job + # permissions this role already holds. + Project.Permissions.RUN_TRAIN_CLASSIFIER_JOB, Project.Permissions.DELETE_OCCURRENCES, Project.Permissions.CREATE_PROJECT_PIPELINE_CONFIG, Project.Permissions.UPDATE_PROJECT_PIPELINE_CONFIG, @@ -154,6 +157,11 @@ class MLDataManager(Role): Project.Permissions.UPDATE_COLLECTION, Project.Permissions.DELETE_COLLECTION, Project.Permissions.POPULATE_COLLECTION, + # Whoever curates the data a model is measured on is the same person who manages + # the data it learns from, so this sits with the collection permissions. + Project.Permissions.CREATE_OCCURRENCE_SET, + Project.Permissions.UPDATE_OCCURRENCE_SET, + Project.Permissions.DELETE_OCCURRENCE_SET, } diff --git a/config/api_router.py b/config/api_router.py index 56a01c591..790acbc53 100644 --- a/config/api_router.py +++ b/config/api_router.py @@ -41,6 +41,7 @@ router.register(r"detections", views.DetectionViewSet) # Register before r"occurrences" — see docs/claude/reference/api-stats-pattern.md (guarded by regression test). router.register(r"occurrences/stats", views.OccurrenceStatsViewSet, basename="occurrence-stats") +router.register(r"occurrences/sets", views.OccurrenceSetViewSet) router.register(r"occurrences", views.OccurrenceViewSet) router.register(r"taxa/lists", views.TaxaListViewSet) # NESTED: /taxa/lists/{taxalist_id}/taxa/ @@ -53,6 +54,7 @@ router.register(r"taxa", views.TaxonViewSet) router.register(r"tags", views.TagViewSet) router.register(r"ml/algorithms", ml_views.AlgorithmViewSet) +router.register(r"ml/training-data", ml_views.TrainingDataViewSet, basename="training-data") router.register(r"ml/labels", ml_views.AlgorithmCategoryMapViewSet) router.register(r"ml/pipelines", ml_views.PipelineViewSet) router.register(r"ml/processing_services", ml_views.ProcessingServiceViewSet) diff --git a/ui/src/components/filtering/filter-control.tsx b/ui/src/components/filtering/filter-control.tsx index e9cc0eed5..552724853 100644 --- a/ui/src/components/filtering/filter-control.tsx +++ b/ui/src/components/filtering/filter-control.tsx @@ -10,6 +10,7 @@ import { DateFilter } from './filters/date-filter' import { DeviceFilter } from './filters/device-filter' import { ImageFilter } from './filters/image-filter' import { OccurrenceFilter } from './filters/occurrence-filter' +import { OccurrenceSetFilter } from './filters/occurrence-set-filter' import { PipelineFilter } from './filters/pipeline-filter' import { SessionFilter } from './filters/session-filter' import { SiteFilter } from './filters/site-filter' @@ -45,6 +46,7 @@ const ComponentMap: { not_algorithm: NotAlgorithmFilter, not_tag_id: TagFilter, not_taxa_list_id: TaxaListFilter, + occurrence_set: OccurrenceSetFilter, pipeline: PipelineFilter, similar_to: OccurrenceFilter, source_image_collection: CaptureSetFilter, diff --git a/ui/src/components/filtering/filters/occurrence-set-filter.tsx b/ui/src/components/filtering/filters/occurrence-set-filter.tsx new file mode 100644 index 000000000..3e6b2c620 --- /dev/null +++ b/ui/src/components/filtering/filters/occurrence-set-filter.tsx @@ -0,0 +1,17 @@ +import { API_ROUTES } from 'data-services/constants' +import { EntityPicker } from 'nova-ui-kit' +import { FilterProps } from './types' + +export const OccurrenceSetFilter = ({ onAdd, onClear, value }: FilterProps) => ( + { + if (value) { + onAdd(value) + } else { + onClear() + } + }} + value={value} + /> +) diff --git a/ui/src/data-services/constants.ts b/ui/src/data-services/constants.ts index 59bae6ede..26891161f 100644 --- a/ui/src/data-services/constants.ts +++ b/ui/src/data-services/constants.ts @@ -3,6 +3,8 @@ export const API_URL = '/api/v2' export const API_ROUTES = { ALGORITHM: 'ml/algorithms', CAPTURE_SET_CHOICES: 'captures/collections/choices', + OCCURRENCE_SETS: 'occurrences/sets', + OCCURRENCE_SET_CHOICES: 'occurrences/sets/choices', CAPTURE_SETS: 'captures/collections', CAPTURES: 'captures', CLASSIFICATIONS: 'classifications', diff --git a/ui/src/data-services/hooks/algorithm/useTrainingSummary.ts b/ui/src/data-services/hooks/algorithm/useTrainingSummary.ts new file mode 100644 index 000000000..db273bf79 --- /dev/null +++ b/ui/src/data-services/hooks/algorithm/useTrainingSummary.ts @@ -0,0 +1,91 @@ +import { FetchParams } from 'data-services/types' +import { getFetchUrl } from 'data-services/utils' +import { useMemo } from 'react' +import { useAuthorizedQuery } from '../auth/useAuthorizedQuery' + +const COLLECTION = 'ml/training-data/summary' + +type ServerTrainingSummary = any // TODO: Update this type + +interface TrainingSummary { + numClasses: number + numOccurrences: number + numRows: number + numTest: number + numTrain: number + numWithoutEmbedding: number +} + +const convertServerRecord = ( + record: ServerTrainingSummary +): TrainingSummary => ({ + numClasses: record.trainable_classes, + numOccurrences: record.occurrences, + numRows: record.rows, + numTest: record.test, + numTrain: record.train, + numWithoutEmbedding: record.verified_detections_without_embedding, +}) + +export const useTrainingSummary = ({ + algorithm, + minPerSpecies, + occurrenceSet, + projectId, + testFraction, +}: { + algorithm?: string + minPerSpecies?: number + occurrenceSet?: string + projectId?: string + testFraction?: number +}): { + summary?: TrainingSummary + isLoading: boolean + isFetching: boolean + error?: unknown +} => { + // A ratio outside (0, 1) is refused by the server, so it is not worth asking. + const ratio = + testFraction !== undefined && testFraction > 0 && testFraction < 1 + ? testFraction + : undefined + + // Settings are sent only when set, so the server answers with the algorithm's own. + const params: FetchParams = { + projectId, + filters: [ + { field: 'algorithm', value: algorithm }, + { field: 'occurrence_set', value: occurrenceSet }, + { + field: 'test_fraction', + value: ratio !== undefined ? `${ratio}` : undefined, + }, + { + field: 'min_per_species', + value: minPerSpecies !== undefined ? `${minPerSpecies}` : undefined, + }, + ], + } + + const fetchUrl = getFetchUrl({ collection: COLLECTION, params }) + + const { data, isLoading, isFetching, error } = + useAuthorizedQuery({ + enabled: !!projectId && !!algorithm, + queryKey: [COLLECTION, params], + url: fetchUrl, + }) + + const summary = useMemo( + () => (data ? convertServerRecord(data) : undefined), + [data] + ) + + return { + summary, + isLoading, + isFetching, + error, + } +} diff --git a/ui/src/data-services/hooks/jobs/useCreateJob.ts b/ui/src/data-services/hooks/jobs/useCreateJob.ts index 4d128a9cd..ef97d5ed0 100644 --- a/ui/src/data-services/hooks/jobs/useCreateJob.ts +++ b/ui/src/data-services/hooks/jobs/useCreateJob.ts @@ -6,7 +6,9 @@ import { useUser } from 'utils/user/userContext' interface JobFieldValues { delay?: number + jobType?: string name: string + params?: Record projectId: string pipeline?: string sourceImage?: string @@ -21,6 +23,10 @@ const convertToServerFieldValues = (fieldValues: JobFieldValues) => ({ pipeline_id: fieldValues.pipeline, source_image_collection_id: fieldValues.sourceImages, source_image_single_id: fieldValues.sourceImage, + // Left out unless set, so the API keeps applying its own defaults: an ML job with no + // type, and no params for a type that reads none. + ...(fieldValues.jobType ? { job_type_key: fieldValues.jobType } : {}), + ...(fieldValues.params ? { params: fieldValues.params } : {}), }) export const useCreateJob = (onSuccess?: (id: string) => void) => { diff --git a/ui/src/data-services/hooks/occurrence-sets/useCreateOccurrenceSet.ts b/ui/src/data-services/hooks/occurrence-sets/useCreateOccurrenceSet.ts new file mode 100644 index 000000000..8d6b6154c --- /dev/null +++ b/ui/src/data-services/hooks/occurrence-sets/useCreateOccurrenceSet.ts @@ -0,0 +1,41 @@ +import { useMutation, useQueryClient } from '@tanstack/react-query' +import axios from 'axios' +import { API_ROUTES, API_URL, SUCCESS_TIMEOUT } from 'data-services/constants' +import { getAuthHeader } from 'data-services/utils' +import { useUser } from 'utils/user/userContext' + +interface OccurrenceSetFieldValues { + name: string + description?: string + projectId: string + occurrenceIds: string[] +} + +const convertToServerFieldValues = (fieldValues: OccurrenceSetFieldValues) => ({ + name: fieldValues.name, + description: fieldValues.description ?? '', + project_id: fieldValues.projectId, + occurrence_ids: fieldValues.occurrenceIds.map((id) => Number(id)), +}) + +export const useCreateOccurrenceSet = (onSuccess?: (id: string) => void) => { + const { user } = useUser() + const queryClient = useQueryClient() + + const { mutateAsync, isLoading, isSuccess, reset, error } = useMutation({ + mutationFn: (fieldValues: OccurrenceSetFieldValues) => + axios.post<{ id: number }>( + `${API_URL}/${API_ROUTES.OCCURRENCE_SETS}/`, + convertToServerFieldValues(fieldValues), + { headers: getAuthHeader(user) } + ), + onSuccess: ({ data }) => { + queryClient.invalidateQueries([API_ROUTES.OCCURRENCE_SETS]) + queryClient.invalidateQueries([API_ROUTES.OCCURRENCE_SET_CHOICES]) + onSuccess?.(`${data.id}`) + setTimeout(reset, SUCCESS_TIMEOUT) + }, + }) + + return { createOccurrenceSet: mutateAsync, isLoading, isSuccess, error } +} diff --git a/ui/src/data-services/models/algorithm.ts b/ui/src/data-services/models/algorithm.ts index c36d507c4..cbe7248e5 100644 --- a/ui/src/data-services/models/algorithm.ts +++ b/ui/src/data-services/models/algorithm.ts @@ -24,6 +24,18 @@ export class Algorithm extends Entity { return this._algorithm.uri } + get trainable(): boolean { + return !!this._algorithm.trainable + } + + /** Settings a retrain would use, which the training job form starts from. */ + get trainingConfig(): { + min_per_species?: number + test_fraction?: number + } { + return this._algorithm.training_config ?? {} + } + get taskType(): string { return snakeCaseToSentenceCase(this._algorithm.task_type) } diff --git a/ui/src/data-services/models/job.ts b/ui/src/data-services/models/job.ts index 8d919e66a..806b66d57 100644 --- a/ui/src/data-services/models/job.ts +++ b/ui/src/data-services/models/job.ts @@ -22,6 +22,7 @@ export const SERVER_JOB_TYPES = [ 'populate_captures_collection', 'data_export', 'post_processing', + 'train_classifier', 'unknown', ] as const @@ -166,6 +167,7 @@ export class Job extends Entity { populate_captures_collection: 'Populate captures collection', data_export: 'Data export', post_processing: 'Post-processing', + train_classifier: 'Train classifier', unknown: 'Unknown', }[key] diff --git a/ui/src/pages/job-details/job-details-form/job-details-form.tsx b/ui/src/pages/job-details/job-details-form/job-details-form.tsx index 8e9eef45e..b2a8a2a93 100644 --- a/ui/src/pages/job-details/job-details-form/job-details-form.tsx +++ b/ui/src/pages/job-details/job-details-form/job-details-form.tsx @@ -16,6 +16,7 @@ import { EntityPicker, InputContent, SaveButton, + Select, } from 'nova-ui-kit' import { CaptureSetPicker } from 'nova-ui-kit/components/select/capture-set-picker' import { useForm } from 'react-hook-form' @@ -23,17 +24,42 @@ import { useParams } from 'react-router-dom' import { APP_ROUTES, DOCS_LINKS } from 'utils/constants' import { STRING, translate } from 'utils/language' import { useFormError } from 'utils/useFormError' - -interface JobFormValues { - delay: number - name: string - pipeline?: string - sourceImage?: string - sourceImages?: string - startNow?: boolean -} +import { TrainClassifierFields } from './train-classifier-fields' +import { JOB_TYPE_ML, JOB_TYPE_TRAIN_CLASSIFIER, JobFormValues } from './types' const config: FormConfig = { + algorithmKey: { + label: translate(STRING.FIELD_LABEL_ALGORITHM), + description: translate(STRING.MESSAGE_TRAINING_ALGORITHM), + rules: { + required: true, + }, + }, + jobType: { + label: translate(STRING.FIELD_LABEL_JOB_TYPE), + rules: { + required: true, + }, + }, + minPerSpecies: { + label: translate(STRING.FIELD_LABEL_MIN_PER_SPECIES), + description: translate(STRING.MESSAGE_MIN_PER_SPECIES), + rules: { + min: 1, + }, + }, + occurrenceSet: { + label: translate(STRING.FIELD_LABEL_OCCURRENCE_SET), + description: translate(STRING.MESSAGE_TRAINING_OCCURRENCE_SET), + }, + testFraction: { + label: translate(STRING.FIELD_LABEL_TEST_FRACTION), + description: translate(STRING.MESSAGE_TEST_FRACTION), + rules: { + min: 0.01, + max: 0.99, + }, + }, delay: { label: translate(STRING.FIELD_LABEL_DELAY), rules: { @@ -82,15 +108,26 @@ export const JobDetailsForm = ({ control, handleSubmit, setError: setFieldError, + watch, } = useForm({ defaultValues: { name: '', delay: 0, + jobType: JOB_TYPE_ML, pipeline: project?.settings.defaultProcessingPipeline?.id, + // Empty rather than unset: a retrain leaves these to the algorithm's own settings. + testFraction: '', + minPerSpecies: '', }, mode: 'onChange', + // Each job type shows its own fields, and the ones hidden must not hold back a + // submission with rules that do not apply to the chosen type. + shouldUnregister: true, }) + const jobType = watch('jobType') + const isTraining = jobType === JOB_TYPE_TRAIN_CLASSIFIER + const errorMessage = useFormError({ error, setFieldError }) return ( @@ -103,91 +140,116 @@ export const JobDetailsForm = ({ /> ) : null} -
- - -
- - - - + {!isTraining ? ( +
+ + +
+ ) : null} ( - )} /> - + + + ( - - - - )} /> + {isTraining ? ( + + ) : ( + + ( + + + + )} + /> + ( + + + + )} + /> + + )} ) } + +export const JobTypePicker = ({ + value, + onValueChange, +}: { + value: string + onValueChange: (value: string) => void +}) => { + const options = [ + { value: JOB_TYPE_ML, label: translate(STRING.JOB_TYPE_ML) }, + { + value: JOB_TYPE_TRAIN_CLASSIFIER, + label: translate(STRING.JOB_TYPE_TRAIN_CLASSIFIER), + }, + ] + + return ( + + + + + + {options.map((option) => ( + + {option.label} + + ))} + + + ) +} diff --git a/ui/src/pages/job-details/job-details-form/train-classifier-fields.tsx b/ui/src/pages/job-details/job-details-form/train-classifier-fields.tsx new file mode 100644 index 000000000..db33838c4 --- /dev/null +++ b/ui/src/pages/job-details/job-details-form/train-classifier-fields.tsx @@ -0,0 +1,156 @@ +import { FormController } from 'components/form/form-controller' +import { FormField } from 'components/form/form-field' +import { FormMessage, FormRow } from 'components/form/layout/layout' +import { FormConfig } from 'components/form/types' +import { API_ROUTES } from 'data-services/constants' +import { useAlgorithms } from 'data-services/hooks/algorithm/useAlgorithms' +import { useTrainingSummary } from 'data-services/hooks/algorithm/useTrainingSummary' +import { EntityPicker, InputContent, Select } from 'nova-ui-kit' +import { Control, useWatch } from 'react-hook-form' +import { useParams } from 'react-router-dom' +import { STRING, translate } from 'utils/language' +import { JobFormValues } from './types' + +const toNumber = (value?: number | string) => { + const parsed = Number(value) + return value === undefined || value === '' || isNaN(parsed) + ? undefined + : parsed +} + +export const TrainClassifierFields = ({ + config, + control, +}: { + config: FormConfig + control: Control +}) => { + const { projectId } = useParams() + const values = useWatch({ control }) + const { summary } = useTrainingSummary({ + algorithm: values.algorithmKey, + minPerSpecies: toNumber(values.minPerSpecies), + occurrenceSet: values.occurrenceSet, + projectId, + testFraction: toNumber(values.testFraction), + }) + + return ( + <> + + ( + + + + )} + /> + ( + + + + )} + /> + + + + + + {summary ? ( + summary.numRows === 0 ? ( + + ) : ( + + ) + ) : null} + {summary && summary.numWithoutEmbedding > 0 ? ( + + ) : null} + + ) +} + +export const TrainableAlgorithmPicker = ({ + value, + onValueChange, +}: { + value?: string + onValueChange: (value: string) => void +}) => { + const { projectId } = useParams() + const { algorithms = [], isLoading } = useAlgorithms({ + projectId: projectId as string, + filters: [{ field: 'trainable', value: 'true' }], + }) + + return ( + + + + + + {algorithms.map((algorithm) => ( + + {algorithm.name} + + ))} + + + ) +} diff --git a/ui/src/pages/job-details/job-details-form/types.ts b/ui/src/pages/job-details/job-details-form/types.ts new file mode 100644 index 000000000..f36803884 --- /dev/null +++ b/ui/src/pages/job-details/job-details-form/types.ts @@ -0,0 +1,16 @@ +export const JOB_TYPE_ML = 'ml' +export const JOB_TYPE_TRAIN_CLASSIFIER = 'train_classifier' + +export interface JobFormValues { + algorithmKey?: string + delay: number + jobType: string + minPerSpecies?: number | string + name: string + occurrenceSet?: string + pipeline?: string + sourceImage?: string + sourceImages?: string + startNow?: boolean + testFraction?: number | string +} diff --git a/ui/src/pages/job-details/new-job-dialog.tsx b/ui/src/pages/job-details/new-job-dialog.tsx index be41fa364..e9698742a 100644 --- a/ui/src/pages/job-details/new-job-dialog.tsx +++ b/ui/src/pages/job-details/new-job-dialog.tsx @@ -5,6 +5,7 @@ import { useState } from 'react' import { useParams } from 'react-router-dom' import { STRING, translate } from 'utils/language' import { JobDetailsForm } from './job-details-form/job-details-form' +import { JOB_TYPE_TRAIN_CLASSIFIER } from './job-details-form/types' import styles from './job-details.module.scss' const CLOSE_TIMEOUT = 1000 @@ -37,9 +38,30 @@ export const NewJobDialog = () => { error={error} isLoading={isLoading} isSuccess={isSuccess} - onSubmit={(data) => { + onSubmit={({ + algorithmKey, + jobType, + minPerSpecies, + occurrenceSet, + testFraction, + ...data + }) => { createJob({ ...data, + jobType, + params: + jobType === JOB_TYPE_TRAIN_CLASSIFIER + ? { + algorithm_key: algorithmKey, + occurrence_set_id: occurrenceSet + ? Number(occurrenceSet) + : undefined, + // A field left empty means the algorithm's own setting, so it + // is left out rather than sent as an empty string. + test_fraction: testFraction || undefined, + min_per_species: minPerSpecies || undefined, + } + : undefined, projectId: projectId as string, }) }} diff --git a/ui/src/pages/occurrences/create-occurrence-set/create-occurrence-set-popover.tsx b/ui/src/pages/occurrences/create-occurrence-set/create-occurrence-set-popover.tsx new file mode 100644 index 000000000..59025acc0 --- /dev/null +++ b/ui/src/pages/occurrences/create-occurrence-set/create-occurrence-set-popover.tsx @@ -0,0 +1,39 @@ +import { BookmarkPlusIcon } from 'lucide-react' +import { Button, Popover } from 'nova-ui-kit' +import { useState } from 'react' +import { STRING, translate } from 'utils/language' +import { CreateOccurrenceSet } from './create-occurrence-set' + +interface CreateOccurrenceSetPopoverProps { + occurrenceIds: string[] +} + +export const CreateOccurrenceSetPopover = ({ + occurrenceIds, +}: CreateOccurrenceSetPopoverProps) => { + const [open, setOpen] = useState(false) + + return ( + + + + + + setOpen(false)} + /> + + + ) +} diff --git a/ui/src/pages/occurrences/create-occurrence-set/create-occurrence-set.tsx b/ui/src/pages/occurrences/create-occurrence-set/create-occurrence-set.tsx new file mode 100644 index 000000000..f249c36f9 --- /dev/null +++ b/ui/src/pages/occurrences/create-occurrence-set/create-occurrence-set.tsx @@ -0,0 +1,70 @@ +import { FormError } from 'components/form/layout/layout' +import { useCreateOccurrenceSet } from 'data-services/hooks/occurrence-sets/useCreateOccurrenceSet' +import { Loader2Icon } from 'lucide-react' +import { Button, Input } from 'nova-ui-kit' +import { useState } from 'react' +import { useParams } from 'react-router-dom' +import { STRING, translate } from 'utils/language' +import { parseServerError } from 'utils/parseServerError/parseServerError' + +interface CreateOccurrenceSetProps { + occurrenceIds: string[] + onCancel: () => void +} + +export const CreateOccurrenceSet = ({ + occurrenceIds, + onCancel, +}: CreateOccurrenceSetProps) => { + const { projectId } = useParams() + const [name, setName] = useState('') + const { createOccurrenceSet, isLoading, error } = + useCreateOccurrenceSet(onCancel) + const formError = error ? parseServerError(error)?.message : undefined + + return ( + <> + {formError && ( + + )} +
+
+ setName(e.target.value)} + value={name} + /> +
+ + {translate(STRING.MESSAGE_OCCURRENCE_SET_IS_FIXED, { + count: `${occurrenceIds.length}`, + })} + +
+ + +
+
+ + ) +} diff --git a/ui/src/pages/occurrences/occurrence-filters.ts b/ui/src/pages/occurrences/occurrence-filters.ts index e63d5093c..6a79c7741 100644 --- a/ui/src/pages/occurrences/occurrence-filters.ts +++ b/ui/src/pages/occurrences/occurrence-filters.ts @@ -11,6 +11,7 @@ export const FILTERS_TO_OCCURRENCES = [ 'verified', 'verified_by_me', 'collection', + 'occurrence_set', 'date_start', 'date_end', 'deployment', diff --git a/ui/src/pages/occurrences/occurrences.tsx b/ui/src/pages/occurrences/occurrences.tsx index bbffc1e12..1fda091c3 100644 --- a/ui/src/pages/occurrences/occurrences.tsx +++ b/ui/src/pages/occurrences/occurrences.tsx @@ -31,6 +31,7 @@ import { columns } from './occurrence-columns' import { OccurrenceDetailsDialog } from './occurrence-details-dialog' import { OccurrenceGallery } from './occurrence-gallery' import { OccurrenceStats } from './occurrence-stats' +import { CreateOccurrenceSetPopover } from './create-occurrence-set/create-occurrence-set-popover' import { OccurrencesActions } from './occurrences-actions' export const Occurrences = () => { @@ -114,6 +115,7 @@ export const Occurrences = () => { defaultOpen={someActive( [ 'collection', + 'occurrence_set', 'deployment', 'deployment__device', 'deployment__research_site', @@ -127,6 +129,7 @@ export const Occurrences = () => { + @@ -167,6 +170,9 @@ export const Occurrences = () => { Export + {selectedItems.length > 0 && ( + + )} [ }, }, }, + { + label: translate(STRING.FIELD_LABEL_OCCURRENCE_SET), + field: 'occurrence_set', + tooltip: { + text: translate(STRING.TOOLTIP_OCCURRENCE_SET), + }, + }, { label: translate(STRING.FIELD_LABEL_CAPTURE_SET), field: 'collection', // This is for viewing occurrences by capture set. @TODO: Can we update this key to "capture_set_id" to streamline?