Skip to content

Commit

Permalink
adding a pipeline get/set to model core class
Browse files Browse the repository at this point in the history
  • Loading branch information
brifordwylie committed Apr 22, 2024
1 parent 913d67c commit 69baf38
Showing 1 changed file with 14 additions and 0 deletions.
14 changes: 14 additions & 0 deletions src/sageworks/core/artifacts/model_core.py
Original file line number Diff line number Diff line change
Expand Up @@ -377,6 +377,7 @@ def details(self, recompute=False) -> dict:

self.log.info("Recomputing Model Details...")
details = self.summary()
details["pipeline"] = self.get_pipeline()
details["model_type"] = self.model_type.value
details["model_package_group_arn"] = self.group_arn()
details["model_package_arn"] = self.model_package_arn()
Expand Down Expand Up @@ -414,6 +415,19 @@ def details(self, recompute=False) -> dict:
# Return the details
return details

# Pipeline for this model
def get_pipeline(self) -> str:
"""Get the pipeline for this model"""
return self.sageworks_meta().get("sageworks_pipeline")

def set_pipeline(self, pipeline: str):
"""Set the pipeline for this model
Args:
pipeline (str): Pipeline that was used to create this model
"""
self.upsert_sageworks_meta({"sageworks_pipeline": pipeline})

def expected_meta(self) -> list[str]:
"""Metadata we expect to see for this Model when it's ready
Returns:
Expand Down

0 comments on commit 69baf38

Please sign in to comment.