Skip to content

Commit f900f6d

Browse files
committed
use datetime_to_ts_millis helper
1 parent a517173 commit f900f6d

File tree

1 file changed

+10
-9
lines changed
  • metadata-ingestion/src/datahub/ingestion/source

1 file changed

+10
-9
lines changed

metadata-ingestion/src/datahub/ingestion/source/vertexai.py

+10-9
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import dataclasses
22
import logging
3-
import time
43
from collections import defaultdict
4+
from datetime import datetime
55
from typing import Any, Iterable, List, Optional, TypeVar
66

77
from google.api_core.exceptions import GoogleAPICallError
@@ -54,6 +54,7 @@
5454
VersionTagClass,
5555
)
5656
from datahub.utilities.str_enum import StrEnum
57+
from datahub.utilities.time import datetime_to_ts_millis
5758

5859
T = TypeVar("T")
5960

@@ -272,9 +273,9 @@ def _gen_training_job_mcps(
272273
job_urn = builder.make_data_process_instance_urn(job_id)
273274

274275
created_time = (
275-
int(job.create_time.timestamp() * 1000)
276+
datetime_to_ts_millis(job.create_time)
276277
if job.create_time
277-
else int(time.time() * 1000)
278+
else datetime_to_ts_millis(datetime.now())
278279
)
279280
created_actor = f"urn:li:platformResource:{self.platform}"
280281

@@ -332,11 +333,11 @@ def _gen_ml_group_mcps(
332333
MLModelGroupPropertiesClass(
333334
name=self._make_vertexai_model_group_name(model.name),
334335
description=model.description,
335-
created=TimeStampClass(time=int(model.create_time.timestamp() * 1000))
336+
created=TimeStampClass(time=datetime_to_ts_millis(model.create_time))
336337
if model.create_time
337338
else None,
338339
lastModified=TimeStampClass(
339-
time=int(model.update_time.timestamp() * 1000)
340+
time=datetime_to_ts_millis(model.update_time)
340341
)
341342
if model.update_time
342343
else None,
@@ -427,7 +428,7 @@ def _get_input_dataset_mcps(
427428
aspects.append(
428429
DatasetPropertiesClass(
429430
name=self._make_vertexai_dataset_name(ds.name),
430-
created=TimeStampClass(time=int(ds.create_time.timestamp() * 1000))
431+
created=TimeStampClass(time=datetime_to_ts_millis(ds.create_time))
431432
if ds.create_time
432433
else None,
433434
description=f"Dataset: {ds.display_name}",
@@ -528,7 +529,7 @@ def _gen_endpoint_mcps(
528529
aspects.append(
529530
MLModelDeploymentPropertiesClass(
530531
description=model.description,
531-
createdAt=int(endpoint.create_time.timestamp() * 1000),
532+
createdAt=datetime_to_ts_millis(endpoint.create_time),
532533
version=VersionTagClass(
533534
versionTag=str(model_version.version_id)
534535
),
@@ -595,12 +596,12 @@ def _gen_ml_model_mcps(
595596
"resourceName": model.resource_name,
596597
},
597598
created=TimeStampClass(
598-
int(model_version.version_create_time.timestamp() * 1000)
599+
datetime_to_ts_millis(model_version.version_create_time)
599600
)
600601
if model_version.version_create_time
601602
else None,
602603
lastModified=TimeStampClass(
603-
int(model_version.version_update_time.timestamp() * 1000)
604+
datetime_to_ts_millis(model_version.version_update_time)
604605
)
605606
if model_version.version_update_time
606607
else None,

0 commit comments

Comments
 (0)