Skip to content

Commit

Permalink
Refactor current model version retrieval to use ModelStages and impro…
Browse files Browse the repository at this point in the history
…ve error handling
  • Loading branch information
safoinme committed Nov 27, 2024
1 parent 0b09eb6 commit e7fe052
Showing 1 changed file with 7 additions and 4 deletions.
11 changes: 7 additions & 4 deletions train_and_deploy/steps/promotion/promote_with_metric_compare.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
from zenml import Model, get_step_context, step
from zenml.enums import ModelStages
from zenml.logger import get_logger

logger = get_logger(__name__)
Expand Down Expand Up @@ -54,10 +55,12 @@ def promote_with_metric_compare(

# Get model version numbers from Model Control Plane
latest_version = get_step_context().model
current_version = Model(name=latest_version.name, version=target_env)

current_version_number = current_version.number

current_version_number = None
try:
current_version = Model(name=latest_version.name, version=ModelStages.STAGING)
current_version_number = current_version.number
except Exception:
pass
if current_version_number is None:
logger.info("No current model version found - promoting latest")
else:
Expand Down

0 comments on commit e7fe052

Please sign in to comment.