AGENT-5867 Fix Custom Models GPU environments to report status (#1055)

datarobot · Jun 13, 2024 · fdcad03 · fdcad03
1 parent 97d5881
commit fdcad03
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 5 deletions.
diff --git a/custom_model_runner/datarobot_drum/drum/gpu_predictors/__init__.py b/custom_model_runner/datarobot_drum/drum/gpu_predictors/__init__.py
@@ -31,12 +31,12 @@ def __init__(
         self.current_stage = 0
 
     def report_deployment(self, message: str):
-        if not self.deployment_id:
+        if not all([self.deployment_id, self.mlops_service_url, self.mlops_api_token]):
             # custom model testing has no deployment_id, skip reporting
             return
 
         self.logger.info(message)
-        remote_events_url = f"{self.mlops_service_url}/api/v2/remoteEvents/"
+        remote_events_url = f"{self.mlops_service_url}/remoteEvents/"
 
         title = "Deployment status"
         if self.total_deployment_stages:

diff --git a/custom_model_runner/datarobot_drum/drum/gpu_predictors/base.py b/custom_model_runner/datarobot_drum/drum/gpu_predictors/base.py
@@ -50,6 +50,11 @@ def __init__(self):
         super().__init__()
         self.logger = logging.getLogger(LOGGER_NAME_PREFIX + "." + __name__)
 
+        # used by status reporter
+        self.datarobot_endpoint = os.environ.get("DATAROBOT_ENDPOINT", None)
+        self.datarobot_api_token = os.environ.get("DATAROBOT_API_TOKEN", None)
+        self.deployment_id = os.environ.get("MLOPS_DEPLOYMENT_ID", None)
+
         # server configuration is set in the Drop-in environment
         self.openai_port = os.environ.get("OPENAI_PORT", "9999")
         self.openai_host = os.environ.get("OPENAI_HOST", "localhost")
@@ -103,9 +108,9 @@ def mlpiper_configure(self, params):
             model_dir=self._code_dir, target_type=self.target_type
         )
         self.status_reporter = MLOpsStatusReporter(
-            mlops_service_url=params["external_webserver_url"],
-            mlops_api_token=params["api_token"],
-            deployment_id=params.get("deployment_id"),
+            mlops_service_url=self.datarobot_endpoint,
+            mlops_api_token=self.datarobot_api_token,
+            deployment_id=self.deployment_id,
             verify_ssl=self.verify_ssl,
             total_deployment_stages=self.num_deployment_stages,
         )