fix: modify the code according to the first review.

Signed-off-by: Electronic-Waste <2690692950@qq.com>

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.

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}"
-            )
+            )
+

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")