This is an automated email from the ASF dual-hosted git repository. skrawcz pushed a commit to branch stefan/fix-unit-tests in repository https://gitbox.apache.org/repos/asf/hamilton.git
commit a4a080203b7d6d22fbe412f898f055979c537116 Author: Stefan Krawczyk <[email protected]> AuthorDate: Mon Dec 29 22:32:30 2025 +1100 Updates pydantic model use to use model_dump This removes use of the deprecated way to dump a pydantic model. --- .../scenario_1/fastapi_server.py | 2 +- ui/backend/server/tests/test_lifecycle/test_projects.py | 2 +- .../server/tests/test_lifecycle/test_run_tracking.py | 6 +++--- ui/backend/server/tests/test_lifecycle/test_templates.py | 4 ++-- ui/backend/server/trackingserver_projects/api.py | 4 ++-- ui/backend/server/trackingserver_run_tracking/api.py | 4 ++-- ui/backend/server/trackingserver_run_tracking/schema.py | 8 ++++---- ui/backend/server/trackingserver_template/api.py | 14 ++++++++++---- ui/sdk/src/hamilton_sdk/tracking/pydantic_stats.py | 2 +- 9 files changed, 26 insertions(+), 20 deletions(-) diff --git a/examples/feature_engineering/feature_engineering_multiple_contexts/scenario_1/fastapi_server.py b/examples/feature_engineering/feature_engineering_multiple_contexts/scenario_1/fastapi_server.py index 296c906d..5860a854 100644 --- a/examples/feature_engineering/feature_engineering_multiple_contexts/scenario_1/fastapi_server.py +++ b/examples/feature_engineering/feature_engineering_multiple_contexts/scenario_1/fastapi_server.py @@ -103,7 +103,7 @@ async def predict_model_version1(request: PredictRequest) -> dict: :return: a dictionary with the prediction value. """ # one liner to quickly create some series from the request. - input_series = pd.DataFrame([request.dict()]).to_dict(orient="series") + input_series = pd.DataFrame([request.model_dump()]).to_dict(orient="series") # create the features -- point here is we're reusing the same code as in the training! # with the ability to provide static values for things like `age_mean` and `age_std_dev`. features = await dr.execute( diff --git a/ui/backend/server/tests/test_lifecycle/test_projects.py b/ui/backend/server/tests/test_lifecycle/test_projects.py index 6916f83a..f49f5cb9 100644 --- a/ui/backend/server/tests/test_lifecycle/test_projects.py +++ b/ui/backend/server/tests/test_lifecycle/test_projects.py @@ -25,7 +25,7 @@ async def _setup_sample_project( ) post_response = await async_client.post( "/api/v1/projects", - data=project_to_create.dict(), + data=project_to_create.model_dump(), content_type="application/json", headers={"test_username": username}, ) diff --git a/ui/backend/server/tests/test_lifecycle/test_run_tracking.py b/ui/backend/server/tests/test_lifecycle/test_run_tracking.py index e4f3b784..46332b9b 100644 --- a/ui/backend/server/tests/test_lifecycle/test_run_tracking.py +++ b/ui/backend/server/tests/test_lifecycle/test_run_tracking.py @@ -23,7 +23,7 @@ async def _setup_dag_template( project_id, *_ = await _setup_sample_project(async_client, username) post_dag_template_response = await async_client.post( f"/api/v1/dag_templates?{urlencode({'project_id': project_id})}", - data=dag_template_to_generate.dict(), + data=dag_template_to_generate.model_dump(), content_type="application/json", headers={"test_username": username}, ) @@ -76,7 +76,7 @@ async def test_create_and_check_dag_template_exists(async_client: AsyncClient, d project_id, *_ = await _setup_sample_project(async_client, username) post_dag_template_response = await async_client.post( f"/api/v1/dag_templates?{urlencode({'project_id': project_id})}", - data=dag_template_to_generate.dict(), + data=dag_template_to_generate.model_dump(), content_type="application/json", headers={"test_username": username}, ) @@ -106,7 +106,7 @@ async def test_create_and_update_empty_dag_run(async_client: AsyncClient, db): ) update_dag_run_response = await async_client.put( f"/api/v1/dag_runs/{run_id}/", - data=dag_run_update.dict(), + data=dag_run_update.model_dump(), content_type="application/json", headers={"test_username": username}, ) diff --git a/ui/backend/server/tests/test_lifecycle/test_templates.py b/ui/backend/server/tests/test_lifecycle/test_templates.py index dbb41b0d..126460f3 100644 --- a/ui/backend/server/tests/test_lifecycle/test_templates.py +++ b/ui/backend/server/tests/test_lifecycle/test_templates.py @@ -89,7 +89,7 @@ async def test_create_and_get_dag_template(async_client: AsyncClient, db): project_id, *_ = await _setup_sample_project(async_client, username) post_dag_template_response = await async_client.post( f"/api/v1/dag_templates?{urlencode({'project_id': project_id})}", - data=dag_template_to_generate.dict(), + data=dag_template_to_generate.model_dump(), content_type="application/json", headers={"test_username": username}, ) @@ -140,7 +140,7 @@ async def test_create_and_get_all_project_dag_templates(async_client: AsyncClien ) post_dag_template_response = await async_client.post( f"/api/v1/dag_templates?{urlencode({'project_id': project_id})}", - data=dag_template_to_generate.dict(), + data=dag_template_to_generate.model_dump(), content_type="application/json", headers={"test_username": username}, ) diff --git a/ui/backend/server/trackingserver_projects/api.py b/ui/backend/server/trackingserver_projects/api.py index 3fef9fb0..7c46f6c4 100644 --- a/ui/backend/server/trackingserver_projects/api.py +++ b/ui/backend/server/trackingserver_projects/api.py @@ -151,7 +151,7 @@ async def get_project_by_id( project_id=project_id, type__in=attribute_types ).all() ] - return ProjectOutWithAttributes(**project_out.dict(), attributes=attributes) + return ProjectOutWithAttributes(**project_out.model_dump(), attributes=attributes) @router.put("/v1/projects/{project_id}", response=ProjectOut, tags=["projects"]) @@ -288,7 +288,7 @@ async def get_projects( ) return [ ProjectOutWithAttributes( - **project.dict(), attributes=project_id_to_attributes.get(project.id, []) + **project.model_dump(), attributes=project_id_to_attributes.get(project.id, []) ) for project in projects ] diff --git a/ui/backend/server/trackingserver_run_tracking/api.py b/ui/backend/server/trackingserver_run_tracking/api.py index 055ecda8..c216c3a6 100644 --- a/ui/backend/server/trackingserver_run_tracking/api.py +++ b/ui/backend/server/trackingserver_run_tracking/api.py @@ -64,7 +64,7 @@ async def create_dag_run(request, dag_template_id: int, dag_run: DAGRunIn) -> DA user, teams = request.auth logger.info(f"Creating DAG run for dag template: {dag_template_id} for user: {user.email}") dag_run_created = await DAGRun.objects.acreate( - **dag_run.dict(), dag_template_id=dag_template_id, launched_by_id=user.id + **dag_run.model_dump(), dag_template_id=dag_template_id, launched_by_id=user.id ) logger.info(f"Created DAG run for dag template: {dag_template_id}") return DAGRunOut.from_orm(dag_run_created) @@ -300,7 +300,7 @@ async def bulk_log( ) logger.info(f"Updated {len(task_updates_to_save)} task updates for dag run: {dag_run_id}") node_attributes_to_create = [ - NodeRunAttribute(**attribute.dict(), dag_run_id=dag_run.id) + NodeRunAttribute(**attribute.model_dump(), dag_run_id=dag_run.id) for attribute in node_run_attributes ] logger.info( diff --git a/ui/backend/server/trackingserver_run_tracking/schema.py b/ui/backend/server/trackingserver_run_tracking/schema.py index ef08d9e6..66a7c756 100644 --- a/ui/backend/server/trackingserver_run_tracking/schema.py +++ b/ui/backend/server/trackingserver_run_tracking/schema.py @@ -50,7 +50,7 @@ class DAGRunOut(ModelSchema): def create_with_username(cls, orm_model: DAGRun) -> "DAGRunOut": return DAGRunOut( **{ - **DAGRunOut.from_orm(orm_model).dict(), + **DAGRunOut.from_orm(orm_model).model_dump(), **{ "username_resolved": ( orm_model.launched_by.email if orm_model.launched_by else None @@ -95,7 +95,7 @@ class NodeRunOutWithAttributes(NodeRunOut): def from_data(cls, node_run: NodeRun, attributes: List[NodeRunAttributeOut]): return NodeRunOutWithAttributes( **{ - **NodeRunOut.from_orm(node_run).dict(), + **NodeRunOut.from_orm(node_run).model_dump(), **{"attributes": attributes, "dag_run_id": node_run.dag_run_id}, } ) @@ -109,7 +109,7 @@ class DAGRunOutWithData(DAGRunOut): def from_data(cls, dag_run: DAGRun, node_runs: List[NodeRunOutWithAttributes]): return DAGRunOutWithData( **{ - **DAGRunOut.from_orm(dag_run).dict(), + **DAGRunOut.from_orm(dag_run).model_dump(), # Not sure why this isn't showing up -- todo, clean this up across the board **{"node_runs": node_runs, "dag_template_id": dag_run.dag_template_id}, } @@ -129,7 +129,7 @@ class NodeRunOutWithExtraData(NodeRunOut, BaseModel): def from_orm(cls, obj, dag_template_id): node_run_out = NodeRunOut.from_orm(obj) return NodeRunOutWithExtraData( - **node_run_out.dict(), + **node_run_out.model_dump(), dag_template_id=dag_template_id, dag_run_id=obj.dag_run_id, ) diff --git a/ui/backend/server/trackingserver_template/api.py b/ui/backend/server/trackingserver_template/api.py index 195de34f..04d45c53 100644 --- a/ui/backend/server/trackingserver_template/api.py +++ b/ui/backend/server/trackingserver_template/api.py @@ -74,7 +74,9 @@ async def create_dag_template( code_log = dag_template.code_log if code_log is not None: logger.info(f"Saving code log for project {project_id} for {user.email}") - code_log_url = await blob_store.write_obj("project" + str(project_id), code_log.dict()) + code_log_url = await blob_store.write_obj( + "project" + str(project_id), code_log.model_dump() + ) logger.info(f"Stored code for project {project_id} for {user.email} at {code_log_url}") code_log_store = blob_store.store() code_log_schema_version = 1 @@ -102,7 +104,7 @@ async def create_dag_template( logger.info(f"Created DAG template for project : {project_id} for {user.email}") code_artifacts_created = await CodeArtifact.objects.abulk_create( [ - CodeArtifact(**code_artifact.dict(), dag_template_id=dag_template_created.id) + CodeArtifact(**code_artifact.model_dump(), dag_template_id=dag_template_created.id) for code_artifact in dag_template.code_artifacts ], ignore_conflicts=False, @@ -115,7 +117,11 @@ async def create_dag_template( for node in dag_template.nodes: node_template = NodeTemplate( - **{key: value for key, value in node.dict().items() if key != "code_artifact_pointers"}, + **{ + key: value + for key, value in node.model_dump().items() + if key != "code_artifact_pointers" + }, dag_template_id=dag_template_created.id, ) node_templates_to_create.append(node_template) @@ -400,7 +406,7 @@ async def get_full_dag_templates(request, dag_template_ids: str) -> List[DAGTemp # TODO -- assert that the blob store matches he one we have available code_log = await blob_store.read_obj(dag_template.code_log_url) return DAGTemplateOutWithData( - **dag_template.dict(), + **dag_template.model_dump(), # TODO -- fix this -- this is due to something weird with the ID names in from_orm code_artifacts=code_artifacts_out, nodes=nodes_out, diff --git a/ui/sdk/src/hamilton_sdk/tracking/pydantic_stats.py b/ui/sdk/src/hamilton_sdk/tracking/pydantic_stats.py index f186e380..249447bc 100644 --- a/ui/sdk/src/hamilton_sdk/tracking/pydantic_stats.py +++ b/ui/sdk/src/hamilton_sdk/tracking/pydantic_stats.py @@ -28,7 +28,7 @@ def compute_stats_pydantic( if hasattr(result, "dump_model"): llm_result = result.dump_model() else: - llm_result = result.dict() + llm_result = result.model_dump() return { "observability_type": "dict", "observability_value": {
