From 01d3e2d2263d85017eed259ca9408701f43c4f06 Mon Sep 17 00:00:00 2001
From: Electronic-Waste <2690692950@qq.com>
Date: Mon, 1 Jul 2024 02:27:09 +0000
Subject: [PATCH] fix: modify the code according to the first review.

Signed-off-by: Electronic-Waste <2690692950@qq.com>
---
 sdk/python/v1beta1/kubeflow/katib/__init__.py          |  2 ++
 .../v1beta1/kubeflow/katib/api/katib_report_metrics.py | 10 +++++-----
 .../v1beta1/kubeflow/katib/constants/constants.py      |  3 +++
 3 files changed, 10 insertions(+), 5 deletions(-)

diff --git a/sdk/python/v1beta1/kubeflow/katib/__init__.py b/sdk/python/v1beta1/kubeflow/katib/__init__.py
index a949c8567e1..7bb4961c420 100644
--- a/sdk/python/v1beta1/kubeflow/katib/__init__.py
+++ b/sdk/python/v1beta1/kubeflow/katib/__init__.py
@@ -71,6 +71,8 @@
 
 # Import Katib API client.
 from kubeflow.katib.api.katib_client import KatibClient
+# Import Katib report metrics functions
+from kubeflow.katib.api.katib_report_metrics import report_metrics
 # Import Katib helper functions.
 import kubeflow.katib.api.search as search
 # Import Katib helper constants.
diff --git a/sdk/python/v1beta1/kubeflow/katib/api/katib_report_metrics.py b/sdk/python/v1beta1/kubeflow/katib/api/katib_report_metrics.py
index 7e326d0f6ee..84b0358e777 100644
--- a/sdk/python/v1beta1/kubeflow/katib/api/katib_report_metrics.py
+++ b/sdk/python/v1beta1/kubeflow/katib/api/katib_report_metrics.py
@@ -13,11 +13,10 @@
 # limitations under the License.
 
 import os
-from datetime import datetime
+from datetime import datetime, timezone
 from typing import Any, Dict
 
 import grpc
-import pytz
 import kubeflow.katib.katib_api_pb2 as katib_api_pb2
 from kubeflow.katib.constants import constants
 from kubeflow.katib.utils import utils
@@ -29,7 +28,7 @@ def report_metrics(
 ):
     """Push Metrics Directly to Katib DB
 
-    [!!!] Trial name should always be passed into Katib Trials as env variable `KATIB_TRIAL_NAME`.
+    Katib always pass Trial name as env variable `KATIB_TRIAL_NAME` to the training container.
 
     Args:
         metrics: Dict of metrics pushed to Katib DB.
@@ -56,7 +55,7 @@ def report_metrics(
     
     with katib_api_pb2.beta_create_DBManager_stub(channel) as client:
         try:
-            timestamp = datetime.now(pytz.UTC).isoformat(timespec="nanoseconds")
+            timestamp = datetime.now(timezone.utc).strftime(constants.RFC3339_FORMAT)
             client.ReportObservationLog(
                 request=katib_api_pb2.ReportObservationLogRequest(
                     trial_name=name,
@@ -75,4 +74,5 @@ def report_metrics(
         except Exception as e:
             raise RuntimeError(
                 f"Unable to push metrics to Katib DB for Trial {namespace}/{name}. Exception: {e}"
-            )
\ No newline at end of file
+            )
+        
\ No newline at end of file
diff --git a/sdk/python/v1beta1/kubeflow/katib/constants/constants.py b/sdk/python/v1beta1/kubeflow/katib/constants/constants.py
index 74aa30e4b26..5bdb9911f76 100644
--- a/sdk/python/v1beta1/kubeflow/katib/constants/constants.py
+++ b/sdk/python/v1beta1/kubeflow/katib/constants/constants.py
@@ -17,6 +17,9 @@
 # How long to wait in seconds for requests to the Kubernetes or gRPC API Server.
 DEFAULT_TIMEOUT = 120
 
+# RFC3339 time format
+RFC3339_FORMAT = "%Y-%m-%dT%H:%M:%SZ"
+
 # Global CRD version
 KATIB_VERSION = os.environ.get("EXPERIMENT_VERSION", "v1beta1")