plexus-python-pulse-demeter 1.1.126__py3-none-any.whl

plexus/pulse/demeter/app/api/dataset_workflow_endpoints.py

@@ -0,0 +1,532 @@
+import datetime
+import uuid as py_uuid
+
+import pydantic as pdt
+import sqlalchemy.exc as sa_exc
+import sqlalchemy.orm as sa_orm
+from fastapi import APIRouter, Depends
+from fastapi.params import Body, Query
+from plexus.common.utils.jsonutils import JsonObject
+from plexus.pulse.common.utils.apiutils import managed_db_session
+from plexus.pulse.common.utils.ormutils import (
+    clone_revision_model_instance,
+    db_create_revision_model,
+    db_update_revision_model,
+    model_name_of,
+)
+
+from plexus.pulse.demeter.app.api import current_dt_clock
+from plexus.pulse.demeter.app.core.database import make_session
+from plexus.pulse.demeter.app.models.registry import JobStatusFlag
+from plexus.pulse.demeter.app.models.registry.dataset import (
+    DatasetPromotion as DatasetPromotion,
+    DatasetRegistration as Registration,
+    ModelAssessment as Assessment,
+    ModelAssessmentPromotion as AssessmentPromotion,
+    ModelEvaluation as Evaluation,
+    ModelExperiment as Experiment,
+    ModelExperimentPromotion as ExperimentPromotion,
+    ModelTraining as Training,
+    RegisteredDataset as Dataset,
+)
+from plexus.pulse.demeter.app.models.registry.dataset import (
+    DatasetPromotionTable as DatasetPromotionTable,
+    DatasetRegistrationTable as RegistrationTable,
+    ModelAssessmentPromotionTable as AssessmentPromotionTable,
+    ModelAssessmentTable as AssessmentTable,
+    ModelEvaluationTable as EvaluationTable,
+    ModelExperimentPromotionTable as ExperimentPromotionTable,
+    ModelExperimentTable as ExperimentTable,
+    ModelTrainingTable as TrainingTable,
+    RegisteredDatasetTable as DatasetTable,
+)
+
+router = APIRouter(prefix="/pulse_demeter/dataset_workflow", tags=["dataset_workflow"])
+
+
+class RegistrationRequest(pdt.BaseModel):
+    uuid: py_uuid.UUID = pdt.Field(default_factory=py_uuid.uuid4)
+    resource_path: str
+    props: JsonObject = None
+    dataset: Dataset
+
+
+class RegistrationTableAndDatasetTable(pdt.BaseModel):
+    registration: RegistrationTable
+    dataset: DatasetTable
+
+
+@router.post("/submit_dataset_registration/", response_model=RegistrationTableAndDatasetTable)
+def submit_dataset_registration(
+    request: RegistrationRequest = Body(),
+    current_dt: datetime.datetime = Depends(current_dt_clock),
+    db: sa_orm.Session = Depends(make_session),
+) -> RegistrationTableAndDatasetTable:
+    with managed_db_session(db):
+        db_dataset = db_create_revision_model(db, DatasetTable, request.dataset, current_dt)
+        registration = Registration(uuid=request.uuid,
+                                    dataset_uuid=db_dataset.uuid,
+                                    resource_path=request.resource_path,
+                                    props=request.props)
+        db_registration = db_create_revision_model(db, RegistrationTable, registration, current_dt)
+
+    return RegistrationTableAndDatasetTable(registration=db_registration, dataset=db_dataset)
+
+
+@router.put("/mark_dataset_registration/", response_model=RegistrationTableAndDatasetTable)
+def mark_dataset_registration(
+    uuid: py_uuid.UUID = Query(),
+    flags: int = Query(0),
+    current_dt: datetime.datetime = Depends(current_dt_clock),
+    db: sa_orm.Session = Depends(make_session),
+) -> RegistrationTableAndDatasetTable:
+    with managed_db_session(db):
+        db_registration = (
+            db
+            .query(RegistrationTable)
+            .where(RegistrationTable.expired_at.is_(None), RegistrationTable.uuid == uuid)
+            .one_or_none()
+        )
+        if db_registration is None:
+            raise sa_exc.NoResultFound(f"'{model_name_of(RegistrationTable)}' with uuid '{uuid}' not found")
+
+        registration = clone_revision_model_instance(RegistrationTable, db_registration, clear_meta_fields=False)
+        if registration.flags != flags:
+            registration.flags = flags
+            db_registration = db_update_revision_model(db,
+                                                       RegistrationTable,
+                                                       registration,
+                                                       registration.record_sqn,
+                                                       current_dt)
+
+        db_dataset = (
+            db
+            .query(DatasetTable)
+            .where(DatasetTable.expired_at.is_(None), DatasetTable.uuid == db_registration.dataset_uuid)
+            .one_or_none()
+        )
+        if db_dataset is None:
+            raise sa_exc.NoResultFound(
+                f"'{model_name_of(DatasetTable)}' with uuid '{db_registration.dataset_uuid}' not found")
+
+    return RegistrationTableAndDatasetTable(registration=db_registration, dataset=db_dataset)
+
+
+@router.put("/finish_dataset_registration/", response_model=RegistrationTableAndDatasetTable)
+def finish_dataset_registration(
+    uuid: py_uuid.UUID = Query(),
+    current_dt: datetime.datetime = Depends(current_dt_clock),
+    db: sa_orm.Session = Depends(make_session),
+) -> RegistrationTableAndDatasetTable:
+    return mark_dataset_registration(uuid=uuid, flags=JobStatusFlag.Completed.value, current_dt=current_dt, db=db)
+
+
+@router.put("/fail_dataset_registration/", response_model=RegistrationTableAndDatasetTable)
+def fail_dataset_registration(
+    uuid: py_uuid.UUID = Query(),
+    current_dt: datetime.datetime = Depends(current_dt_clock),
+    db: sa_orm.Session = Depends(make_session),
+) -> RegistrationTableAndDatasetTable:
+    return mark_dataset_registration(uuid=uuid, flags=JobStatusFlag.Failed.value, current_dt=current_dt, db=db)
+
+
+@router.post("/submit_dataset_promotion/", response_model=DatasetPromotionTable)
+def submit_dataset_promotion(
+    uuid: py_uuid.UUID = Query(default_factory=py_uuid.uuid4),
+    dataset_uuid: py_uuid.UUID = Query(),
+    reviewer_sso_sid: int = Query(),
+    props: JsonObject | None = Query(default_factory=dict),
+    current_dt: datetime.datetime = Depends(current_dt_clock),
+    db: sa_orm.Session = Depends(make_session),
+) -> DatasetPromotionTable:
+    with managed_db_session(db):
+        promotion = DatasetPromotion(
+            uuid=uuid,
+            dataset_uuid=dataset_uuid,
+            reviewer_sso_sid=reviewer_sso_sid,
+            props=props,
+        )
+        return db_create_revision_model(db, DatasetPromotionTable, promotion, current_dt)
+
+
+@router.put("/mark_dataset_promotion/", response_model=DatasetPromotionTable)
+def mark_dataset_promotion(
+    uuid: py_uuid.UUID = Query(),
+    flags: int = Query(0),
+    current_dt: datetime.datetime = Depends(current_dt_clock),
+    db: sa_orm.Session = Depends(make_session),
+) -> DatasetPromotionTable:
+    with managed_db_session(db):
+        db_promotion = (
+            db
+            .query(DatasetPromotionTable)
+            .where(DatasetPromotionTable.expired_at.is_(None), DatasetPromotionTable.uuid == uuid)
+            .one_or_none()
+        )
+        if db_promotion is None:
+            raise sa_exc.NoResultFound(f"'{model_name_of(DatasetPromotionTable)}' with uuid '{uuid}' not found")
+
+        promotion = clone_revision_model_instance(DatasetPromotionTable, db_promotion, clear_meta_fields=False)
+        if promotion.flags != flags:
+            promotion.flags = flags
+            db_promotion = db_update_revision_model(db,
+                                                    DatasetPromotionTable,
+                                                    promotion,
+                                                    promotion.record_sqn,
+                                                    current_dt)
+
+    return db_promotion
+
+
+@router.put("/finish_dataset_promotion/", response_model=DatasetPromotionTable)
+def finish_dataset_promotion(
+    uuid: py_uuid.UUID = Query(),
+    current_dt: datetime.datetime = Depends(current_dt_clock),
+    db: sa_orm.Session = Depends(make_session),
+) -> DatasetPromotionTable:
+    return mark_dataset_promotion(uuid=uuid, flags=JobStatusFlag.Completed.value, current_dt=current_dt, db=db)
+
+
+@router.put("/fail_dataset_promotion/", response_model=DatasetPromotionTable)
+def fail_dataset_promotion(
+    uuid: py_uuid.UUID = Query(),
+    current_dt: datetime.datetime = Depends(current_dt_clock),
+    db: sa_orm.Session = Depends(make_session),
+) -> DatasetPromotionTable:
+    return mark_dataset_promotion(uuid=uuid, flags=JobStatusFlag.Failed.value, current_dt=current_dt, db=db)
+
+
+class TrainingRequest(pdt.BaseModel):
+    uuid: py_uuid.UUID = pdt.Field(default_factory=py_uuid.uuid4)
+    props: JsonObject = None
+    experiment: Experiment
+
+
+class TrainingTableAndExperimentTable(pdt.BaseModel):
+    training: TrainingTable
+    experiment: ExperimentTable
+
+
+@router.post("/submit_model_training/", response_model=TrainingTableAndExperimentTable)
+def submit_model_training(
+    request: TrainingRequest = Body(),
+    current_dt: datetime.datetime = Depends(current_dt_clock),
+    db: sa_orm.Session = Depends(make_session),
+) -> TrainingTableAndExperimentTable:
+    with managed_db_session(db):
+        db_experiment = db_create_revision_model(db, ExperimentTable, request.experiment, current_dt)
+        training = Training(uuid=request.uuid, experiment_uuid=db_experiment.uuid, props=request.props)
+        db_training = db_create_revision_model(db, TrainingTable, training, current_dt)
+
+    return TrainingTableAndExperimentTable(training=db_training, experiment=db_experiment)
+
+
+@router.put("/mark_model_training/", response_model=TrainingTableAndExperimentTable)
+def mark_model_training(
+    uuid: py_uuid.UUID = Query(),
+    flags: int = Query(0),
+    current_dt: datetime.datetime = Depends(current_dt_clock),
+    db: sa_orm.Session = Depends(make_session),
+) -> TrainingTableAndExperimentTable:
+    with managed_db_session(db):
+        db_training = (
+            db
+            .query(TrainingTable)
+            .where(TrainingTable.expired_at.is_(None), TrainingTable.uuid == uuid)
+            .one_or_none()
+        )
+        if db_training is None:
+            raise sa_exc.NoResultFound(f"'{model_name_of(TrainingTable)}' with uuid '{uuid}' not found")
+
+        training = clone_revision_model_instance(TrainingTable, db_training, clear_meta_fields=False)
+        if training.flags != flags:
+            training.flags = flags
+            db_training = db_update_revision_model(db, TrainingTable, training, training.record_sqn, current_dt)
+
+        db_experiment = (
+            db
+            .query(ExperimentTable)
+            .where(ExperimentTable.expired_at.is_(None), ExperimentTable.uuid == db_training.experiment_uuid)
+            .one_or_none()
+        )
+        if db_experiment is None:
+            raise sa_exc.NoResultFound(
+                f"'{model_name_of(ExperimentTable)}' with uuid '{db_training.experiment_uuid}' not found")
+
+    return TrainingTableAndExperimentTable(training=db_training, experiment=db_experiment)
+
+
+class ExperimentResult(pdt.BaseModel):
+    path: str
+    checkpoint_path: str
+
+
+@router.put("/finish_model_training/", response_model=TrainingTableAndExperimentTable)
+def finish_model_training(
+    uuid: py_uuid.UUID = Query(),
+    result: ExperimentResult = Body(),
+    current_dt: datetime.datetime = Depends(current_dt_clock),
+    db: sa_orm.Session = Depends(make_session),
+) -> TrainingTableAndExperimentTable:
+    with managed_db_session(db):
+        training_and_experiment = mark_model_training(uuid=uuid,
+                                                      flags=JobStatusFlag.Completed.value,
+                                                      current_dt=current_dt,
+                                                      db=db)
+
+        experiment = clone_revision_model_instance(ExperimentTable,
+                                                   training_and_experiment.experiment,
+                                                   clear_meta_fields=False)
+        experiment.path = result.path
+        experiment.checkpoint_path = result.checkpoint_path
+        db_experiment = db_update_revision_model(db,
+                                                 ExperimentTable,
+                                                 experiment,
+                                                 experiment.record_sqn,
+                                                 current_dt)
+
+    return TrainingTableAndExperimentTable(training=training_and_experiment.training,
+                                           experiment=db_experiment)
+
+
+@router.put("/fail_model_training/", response_model=TrainingTableAndExperimentTable)
+def fail_model_training(
+    uuid: py_uuid.UUID = Query(),
+    current_dt: datetime.datetime = Depends(current_dt_clock),
+    db: sa_orm.Session = Depends(make_session),
+) -> TrainingTableAndExperimentTable:
+    return mark_model_training(uuid=uuid, flags=JobStatusFlag.Failed.value, current_dt=current_dt, db=db)
+
+
+@router.post("/submit_experiment_promotion/", response_model=ExperimentPromotionTable)
+def submit_experiment_promotion(
+    uuid: py_uuid.UUID = Query(default_factory=py_uuid.uuid4),
+    experiment_uuid: py_uuid.UUID = Query(),
+    reviewer_sso_sid: int = Query(),
+    props: JsonObject | None = Query(default_factory=dict),
+    current_dt: datetime.datetime = Depends(current_dt_clock),
+    db: sa_orm.Session = Depends(make_session),
+) -> ExperimentPromotionTable:
+    with managed_db_session(db):
+        promotion = ExperimentPromotion(
+            uuid=uuid,
+            experiment_uuid=experiment_uuid,
+            reviewer_sso_sid=reviewer_sso_sid,
+            props=props,
+        )
+        return db_create_revision_model(db, ExperimentPromotionTable, promotion, current_dt)
+
+
+@router.put("/mark_experiment_promotion/", response_model=ExperimentPromotionTable)
+def mark_experiment_promotion(
+    uuid: py_uuid.UUID = Query(),
+    flags: int = Query(0),
+    current_dt: datetime.datetime = Depends(current_dt_clock),
+    db: sa_orm.Session = Depends(make_session),
+) -> ExperimentPromotionTable:
+    with managed_db_session(db):
+        db_promotion = (
+            db
+            .query(ExperimentPromotionTable)
+            .where(ExperimentPromotionTable.expired_at.is_(None), ExperimentPromotionTable.uuid == uuid)
+            .one_or_none()
+        )
+        if db_promotion is None:
+            raise sa_exc.NoResultFound(f"'{model_name_of(ExperimentPromotionTable)}' with uuid '{uuid}' not found")
+
+        promotion = clone_revision_model_instance(ExperimentPromotionTable, db_promotion, clear_meta_fields=False)
+        if promotion.flags != flags:
+            promotion.flags = flags
+            db_promotion = db_update_revision_model(db,
+                                                    ExperimentPromotionTable,
+                                                    promotion,
+                                                    promotion.record_sqn,
+                                                    current_dt)
+
+    return db_promotion
+
+
+@router.put("/finish_experiment_promotion/", response_model=ExperimentPromotionTable)
+def finish_experiment_promotion(
+    uuid: py_uuid.UUID = Query(),
+    current_dt: datetime.datetime = Depends(current_dt_clock),
+    db: sa_orm.Session = Depends(make_session),
+) -> ExperimentPromotionTable:
+    return mark_experiment_promotion(uuid=uuid, flags=JobStatusFlag.Completed.value, current_dt=current_dt, db=db)
+
+
+@router.put("/fail_experiment_promotion/", response_model=ExperimentPromotionTable)
+def fail_experiment_promotion(
+    uuid: py_uuid.UUID = Query(),
+    current_dt: datetime.datetime = Depends(current_dt_clock),
+    db: sa_orm.Session = Depends(make_session),
+) -> ExperimentPromotionTable:
+    return mark_experiment_promotion(uuid=uuid, flags=JobStatusFlag.Failed.value, current_dt=current_dt, db=db)
+
+
+class EvaluationRequest(pdt.BaseModel):
+    uuid: py_uuid.UUID = pdt.Field(default_factory=py_uuid.uuid4)
+    props: JsonObject = None
+    assessment: Assessment
+
+
+class EvaluationTableAndAssessmentTable(pdt.BaseModel):
+    evaluation: EvaluationTable
+    assessment: AssessmentTable
+
+
+@router.post("/submit_model_evaluation/", response_model=EvaluationTableAndAssessmentTable)
+def submit_model_evaluation(
+    request: EvaluationRequest = Body(),
+    current_dt: datetime.datetime = Depends(current_dt_clock),
+    db: sa_orm.Session = Depends(make_session),
+) -> EvaluationTableAndAssessmentTable:
+    with managed_db_session(db):
+        db_assessment = db_create_revision_model(db, AssessmentTable, request.assessment, current_dt)
+        evaluation = Evaluation(uuid=request.uuid, assessment_uuid=db_assessment.uuid, props=request.props)
+        db_evaluation = db_create_revision_model(db, EvaluationTable, evaluation, current_dt)
+
+    return EvaluationTableAndAssessmentTable(evaluation=db_evaluation, assessment=db_assessment)
+
+
+@router.put("/mark_model_evaluation/", response_model=EvaluationTableAndAssessmentTable)
+def mark_model_evaluation(
+    uuid: py_uuid.UUID = Query(),
+    flags: int = Query(0),
+    current_dt: datetime.datetime = Depends(current_dt_clock),
+    db: sa_orm.Session = Depends(make_session),
+) -> EvaluationTableAndAssessmentTable:
+    with managed_db_session(db):
+        db_evaluation = (
+            db
+            .query(EvaluationTable)
+            .where(EvaluationTable.expired_at.is_(None), EvaluationTable.uuid == uuid)
+            .one_or_none()
+        )
+        if db_evaluation is None:
+            raise sa_exc.NoResultFound(f"'{model_name_of(EvaluationTable)}' with uuid '{uuid}' not found")
+
+        evaluation = clone_revision_model_instance(EvaluationTable, db_evaluation, clear_meta_fields=False)
+        if evaluation.flags != flags:
+            evaluation.flags = flags
+            db_evaluation = db_update_revision_model(db, EvaluationTable, evaluation, evaluation.record_sqn, current_dt)
+
+        db_assessment = (
+            db
+            .query(AssessmentTable)
+            .where(AssessmentTable.expired_at.is_(None), AssessmentTable.uuid == db_evaluation.assessment_uuid)
+            .one_or_none()
+        )
+        if db_assessment is None:
+            raise sa_exc.NoResultFound(
+                f"'{model_name_of(AssessmentTable)}' with uuid '{db_evaluation.assessment_uuid}' not found")
+
+    return EvaluationTableAndAssessmentTable(evaluation=db_evaluation, assessment=db_assessment)
+
+
+class AssessmentResult(pdt.BaseModel):
+    path: str
+    assessment_data: JsonObject
+
+
+@router.put("/finish_model_evaluation/", response_model=EvaluationTableAndAssessmentTable)
+def finish_model_evaluation(
+    uuid: py_uuid.UUID = Query(),
+    result: AssessmentResult = Body(),
+    current_dt: datetime.datetime = Depends(current_dt_clock),
+    db: sa_orm.Session = Depends(make_session),
+) -> EvaluationTableAndAssessmentTable:
+    with managed_db_session(db):
+        evaluation_and_assessment = mark_model_evaluation(uuid=uuid,
+                                                          flags=JobStatusFlag.Completed.value,
+                                                          current_dt=current_dt,
+                                                          db=db)
+
+        assessment = clone_revision_model_instance(AssessmentTable,
+                                                   evaluation_and_assessment.assessment,
+                                                   clear_meta_fields=False)
+        assessment.path = result.path
+        assessment.assessment_data = result.assessment_data
+        db_assessment = db_update_revision_model(db,
+                                                 AssessmentTable,
+                                                 assessment,
+                                                 assessment.record_sqn,
+                                                 current_dt)
+
+    return EvaluationTableAndAssessmentTable(evaluation=evaluation_and_assessment.evaluation,
+                                             assessment=db_assessment)
+
+
+@router.put("/fail_model_evaluation/", response_model=EvaluationTableAndAssessmentTable)
+def fail_model_evaluation(
+    uuid: py_uuid.UUID = Query(),
+    current_dt: datetime.datetime = Depends(current_dt_clock),
+    db: sa_orm.Session = Depends(make_session),
+) -> EvaluationTableAndAssessmentTable:
+    return mark_model_evaluation(uuid=uuid, flags=JobStatusFlag.Failed.value, current_dt=current_dt, db=db)
+
+
+@router.post("/submit_assessment_promotion/", response_model=AssessmentPromotionTable)
+def submit_assessment_promotion(
+    uuid: py_uuid.UUID = Query(default_factory=py_uuid.uuid4),
+    assessment_uuid: py_uuid.UUID = Query(),
+    reviewer_sso_sid: int = Query(),
+    props: JsonObject | None = Query(default_factory=dict),
+    current_dt: datetime.datetime = Depends(current_dt_clock),
+    db: sa_orm.Session = Depends(make_session),
+) -> AssessmentPromotionTable:
+    with managed_db_session(db):
+        promotion = AssessmentPromotion(
+            uuid=uuid,
+            assessment_uuid=assessment_uuid,
+            reviewer_sso_sid=reviewer_sso_sid,
+            props=props,
+        )
+        return db_create_revision_model(db, AssessmentPromotionTable, promotion, current_dt)
+
+
+@router.put("/mark_assessment_promotion/", response_model=AssessmentPromotionTable)
+def mark_assessment_promotion(
+    uuid: py_uuid.UUID = Query(),
+    flags: int = Query(0),
+    current_dt: datetime.datetime = Depends(current_dt_clock),
+    db: sa_orm.Session = Depends(make_session),
+) -> AssessmentPromotionTable:
+    with managed_db_session(db):
+        db_promotion = (
+            db
+            .query(AssessmentPromotionTable)
+            .where(AssessmentPromotionTable.expired_at.is_(None), AssessmentPromotionTable.uuid == uuid)
+            .one_or_none()
+        )
+        if db_promotion is None:
+            raise sa_exc.NoResultFound(f"'{model_name_of(AssessmentPromotionTable)}' with uuid '{uuid}' not found")
+
+        promotion = clone_revision_model_instance(AssessmentPromotionTable, db_promotion, clear_meta_fields=False)
+        if promotion.flags != flags:
+            promotion.flags = flags
+            db_promotion = db_update_revision_model(db,
+                                                    AssessmentPromotionTable,
+                                                    promotion,
+                                                    promotion.record_sqn,
+                                                    current_dt)
+
+    return db_promotion
+
+
+@router.put("/finish_assessment_promotion/", response_model=AssessmentPromotionTable)
+def finish_assessment_promotion(
+    uuid: py_uuid.UUID = Query(),
+    current_dt: datetime.datetime = Depends(current_dt_clock),
+    db: sa_orm.Session = Depends(make_session),
+) -> AssessmentPromotionTable:
+    return mark_assessment_promotion(uuid=uuid, flags=JobStatusFlag.Completed.value, current_dt=current_dt, db=db)
+
+
+@router.put("/fail_assessment_promotion/", response_model=AssessmentPromotionTable)
+def fail_assessment_promotion(
+    uuid: py_uuid.UUID = Query(),
+    current_dt: datetime.datetime = Depends(current_dt_clock),
+    db: sa_orm.Session = Depends(make_session),
+) -> AssessmentPromotionTable:
+    return mark_assessment_promotion(uuid=uuid, flags=JobStatusFlag.Failed.value, current_dt=current_dt, db=db)