diff --git a/frontend/src/api/modelRegistry/__tests__/custom.spec.ts b/frontend/src/api/modelRegistry/__tests__/custom.spec.ts index c3951a6ac6..7388b123bf 100644 --- a/frontend/src/api/modelRegistry/__tests__/custom.spec.ts +++ b/frontend/src/api/modelRegistry/__tests__/custom.spec.ts @@ -118,7 +118,7 @@ describe('createModelVersionForRegisteredModel', () => { state: ModelState.LIVE, customProperties: {}, }), - ).toBe(mockResultPromise); + ).toEqual(mockResultPromise); expect(proxyCREATEMock).toHaveBeenCalledTimes(1); expect(proxyCREATEMock).toHaveBeenCalledWith( 'hostPath', diff --git a/frontend/src/api/modelRegistry/custom.ts b/frontend/src/api/modelRegistry/custom.ts index 5abf055af3..a4761192ff 100644 --- a/frontend/src/api/modelRegistry/custom.ts +++ b/frontend/src/api/modelRegistry/custom.ts @@ -8,11 +8,11 @@ import { ModelVersion, RegisteredModelList, RegisteredModel, - ModelRegistryMetadataType, } from '~/concepts/modelRegistry/types'; +import { MODEL_REGISTRY_API_VERSION } from '~/concepts/modelRegistry/const'; +import { bumpRegisteredModelTimestamp } from '~/concepts/modelRegistry/utils/updateTimestamps'; import { proxyCREATE, proxyGET, proxyPATCH } from '~/api/proxyUtils'; import { K8sAPIOptions } from '~/k8sTypes'; -import { MODEL_REGISTRY_API_VERSION } from '~/concepts/modelRegistry/const'; import { handleModelRegistryFailures } from './errorUtils'; export const createRegisteredModel = @@ -47,7 +47,6 @@ export const createModelVersionForRegisteredModel = registeredModelId: string, data: CreateModelVersionData, ): Promise => { - // First create the version const newVersion = await handleModelRegistryFailures( proxyCREATE( hostpath, @@ -58,22 +57,20 @@ export const createModelVersionForRegisteredModel = ), ); - const currentTime = new Date().toISOString(); - await handleModelRegistryFailures( - proxyPATCH( - hostpath, - `/api/model_registry/${MODEL_REGISTRY_API_VERSION}/registered_models/${registeredModelId}`, - { - state: 'LIVE', - customProperties: { - _lastModified: { - metadataType: ModelRegistryMetadataType.STRING, - stringValue: currentTime, - }, - }, - }, - opts, - ), + // Use the established timestamp update utility + await bumpRegisteredModelTimestamp( + { + patchRegisteredModel: (apiOpts, patchData, id) => + handleModelRegistryFailures( + proxyPATCH( + hostpath, + `/api/model_registry/${MODEL_REGISTRY_API_VERSION}/registered_models/${id}`, + patchData, + apiOpts, + ), + ), + }, + registeredModelId, ); return newVersion; diff --git a/frontend/src/concepts/modelRegistry/utils/updateTimestamps.ts b/frontend/src/concepts/modelRegistry/utils/updateTimestamps.ts index c1e594b3fe..7bc69d564c 100644 --- a/frontend/src/concepts/modelRegistry/utils/updateTimestamps.ts +++ b/frontend/src/concepts/modelRegistry/utils/updateTimestamps.ts @@ -4,6 +4,8 @@ import { ModelRegistryMetadataType, } from '~/concepts/modelRegistry/types'; +type MinimalModelRegistryAPI = Pick; + export const bumpModelVersionTimestamp = async ( api: ModelRegistryAPIs, modelVersionId: string, @@ -24,7 +26,7 @@ export const bumpModelVersionTimestamp = async ( }; export const bumpRegisteredModelTimestamp = async ( - api: ModelRegistryAPIs, + api: MinimalModelRegistryAPI, registeredModelId: string, ): Promise => { if (!registeredModelId) {