Task/dss 12.4.0 bump dataikuapi (#287)

* Sync dataikuapi sources for 12.4.0

* 12.4.0 metadata

Author: kevanescence

Diff for: HISTORY.txt

@@ -2,6 +2,11 @@ Changelog
 ==========
 
 
+12.4.0 (2023-12-06)
+---------------------
+
+* Initial release for DSS 12.4.0
+
 12.3.2 (2023-11-16)
 ---------------------
 

Diff for: dataikuapi/apinode_admin_client.py

@@ -72,3 +72,17 @@ def import_code_env_in_cache(self, file_dir, language):
             "fileDir": file_dir,
             "language": language
         })
+
+    def register_code_env_in_cache(self, exported_env_dir, built_env_dir, language):
+        """
+        Import a code env in global cache from an exported code env base folder
+
+        :param exported_env_dir: path of an exported code env base folder
+        :param built_env_dir: path where the code env was built and is available
+        :param language: language of the code env (`python` or `R`)
+        """
+        return self._perform_json("POST", "register-global-code-env", params={
+            "exportedEnvDir": exported_env_dir,
+            "builtEnvDir": built_env_dir,
+            "language": language
+        })

Diff for: dataikuapi/apinode_client.py

@@ -1,6 +1,7 @@
 from .utils import DataikuException
 from .base_client import DSSBaseClient
 
+
 class APINodeClient(DSSBaseClient):
     """Entry point for the DSS API Node client
     This is an API client for the user-facing API of DSS API Node server (user facing API)
@@ -17,6 +18,13 @@ def __init__(self, uri, service_id, api_key=None, bearer_token=None):
         """
         DSSBaseClient.__init__(self, "%s/%s" % (uri, "public/api/v1/%s" % service_id), api_key=api_key, bearer_token=bearer_token)
 
+    @staticmethod
+    def _set_dispatch(obj, forced_generation, dispatch_key):
+        if forced_generation is not None:
+            obj["dispatch"] = {"forcedGeneration": forced_generation}
+        elif dispatch_key is not None:
+            obj["dispatch"] = {"dispatchKey": dispatch_key}
+
     def predict_record(self, endpoint_id, features, forced_generation=None, dispatch_key=None, context=None,
                        with_explanations=None, explanation_method=None, n_explanations=None, n_explanations_mc_steps=None):
         """
@@ -45,12 +53,9 @@ def predict_record(self, endpoint_id, features, forced_generation=None, dispatch
             }
         }
 
+        self._set_dispatch(obj, forced_generation, dispatch_key)
         if context is not None:
             obj["context"] = context
-        if forced_generation is not None:
-            obj["dispatch"] = {"forcedGeneration" : forced_generation }
-        elif dispatch_key is not None:
-            obj["dispatch"] = {"dispatchKey" : dispatch_key }
 
         return self._perform_json("POST", "%s/predict" % endpoint_id, body = obj)
 
@@ -86,10 +91,7 @@ def predict_records(self, endpoint_id, records, forced_generation=None, dispatch
             }
         }
 
-        if forced_generation is not None:
-            obj["dispatch"] = {"forcedGeneration" : forced_generation }
-        elif dispatch_key is not None:
-            obj["dispatch"] = {"dispatchKey" : dispatch_key }
+        self._set_dispatch(obj, forced_generation, dispatch_key)
 
         return self._perform_json("POST", "%s/predict-multi" % endpoint_id, body = obj)
 
@@ -145,13 +147,52 @@ def forecast(self, endpoint_id, records, forced_generation=None, dispatch_key=No
 
         obj = {"items": records}
 
-        if forced_generation is not None:
-            obj["dispatch"] = {"forcedGeneration": forced_generation}
-        elif dispatch_key is not None:
-            obj["dispatch"] = {"dispatchKey": dispatch_key}
+        self._set_dispatch(obj, forced_generation, dispatch_key)
 
         return self._perform_json("POST", "{}/forecast".format(endpoint_id), body=obj)
 
+
+    def predict_effect(self, endpoint_id, features, forced_generation=None, dispatch_key=None):
+        """
+        Predicts the treatment effect of a single record on a DSS API node endpoint (standard causal prediction)
+
+        :param str endpoint_id: Identifier of the endpoint to query
+        :param features: Python dictionary of features of the record
+        :param forced_generation: See documentation about multi-version prediction
+        :param dispatch_key: See documentation about multi-version prediction
+
+        :return: a Python dict of the API answer. The answer contains a "result" key (itself a dict)
+        """
+        obj = {
+            "features": features,
+        }
+        self._set_dispatch(obj, forced_generation, dispatch_key)
+
+        return self._perform_json("POST", "%s/predict-effect" % endpoint_id, body=obj)
+
+    def predict_effects(self, endpoint_id, records, forced_generation=None, dispatch_key=None):
+        """
+        Predicts the treatment effects on a batch of records on a DSS API node endpoint (standard causal prediction)
+
+        :param str endpoint_id: Identifier of the endpoint to query
+        :param records: Python list of records. Each record must be a Python dict. Each record must contain a "features" dict (see predict_record) and optionally a "context" dict.
+        :param dispatch_key: See documentation about multi-version prediction
+
+        :return: a Python dict of the API answer. The answer contains a "results" key (which is an array of result objects)
+        """
+
+        for record in records:
+            if not "features" in record:
+                raise ValueError("Each record must contain a 'features' dict")
+
+        obj = {
+            "items": records,
+        }
+
+        self._set_dispatch(obj, forced_generation, dispatch_key)
+
+        return self._perform_json("POST", "%s/predict-effect-multi" % endpoint_id, body = obj)
+
     def sql_query(self, endpoint_id, parameters):
         """
         Queries a "SQL query" endpoint on a DSS API node

Diff for: dataikuapi/dss/admin.py

@@ -1,6 +1,7 @@
 from .future import DSSFuture
 import json, warnings
 from datetime import datetime
+from ..utils import _timestamp_ms_to_zoned_datetime
 
 
 class DSSConnectionListItem(dict):
@@ -86,6 +87,15 @@ def get_params(self):
         """
         return self["params"]
 
+    def get_resolved_params(self):
+        """
+        Get the resolved parameters of the connection, as a dict. May be null depending on the connection type.
+
+        :return: the resolved parameters, as a dict. Each connection type has different sets of fields.
+        :rtype: dict
+        """
+        return self["resolvedParams"]
+
     def get_basic_credential(self):
         """
         Get the basic credential (user/password pair) for this connection, if available
@@ -835,7 +845,7 @@ def creation_date(self):
         :rtype: :class:`datetime.datetime` or None
         """
         timestamp = self.settings["creationDate"] if "creationDate" in self.settings else None
-        return datetime.fromtimestamp(timestamp / 1000) if timestamp else None
+        return _timestamp_ms_to_zoned_datetime(timestamp)
 
     def save(self):
         """
@@ -902,7 +912,7 @@ def last_successful_login(self):
         :rtype: :class:`datetime.datetime` or None
         """
         timestamp = self.activity["lastSuccessfulLogin"]
-        return datetime.fromtimestamp(timestamp / 1000) if timestamp > 0 else None
+        return _timestamp_ms_to_zoned_datetime(timestamp)
 
     @property
     def last_failed_login(self):
@@ -915,7 +925,7 @@ def last_failed_login(self):
         :rtype: :class:`datetime.datetime` or None
         """
         timestamp = self.activity["lastFailedLogin"]
-        return datetime.fromtimestamp(timestamp / 1000) if timestamp > 0 else None
+        return _timestamp_ms_to_zoned_datetime(timestamp)
 
     @property
     def last_session_activity(self):
@@ -931,7 +941,7 @@ def last_session_activity(self):
         :rtype: :class:`datetime.datetime` or None
         """
         timestamp = self.activity["lastSessionActivity"]
-        return datetime.fromtimestamp(timestamp / 1000) if timestamp > 0 else None
+        return _timestamp_ms_to_zoned_datetime(timestamp)
 
 
 class DSSAuthorizationMatrix(object):
@@ -1916,7 +1926,7 @@ def created_on(self):
         :rtype: :class:`datetime.datetime`
         """
         timestamp = self["createdOn"]
-        return datetime.fromtimestamp(timestamp / 1000) if timestamp > 0 else None
+        return _timestamp_ms_to_zoned_datetime(timestamp)
 
     @property
     def created_by(self):
@@ -2040,7 +2050,7 @@ def created_on(self):
         :rtype: :class:`datetime.datetime`
         """
         timestamp = self["createdOn"]
-        return datetime.fromtimestamp(timestamp / 1000) if timestamp > 0 else None
+        return _timestamp_ms_to_zoned_datetime(timestamp)
 
     @property
     def created_by(self):
@@ -2514,10 +2524,7 @@ def last_built(self):
         :rtype: :class:`datetime.datetime`
         """
         ts = self._data.get("lastBuilt", 0)
-        if ts > 0:
-            return datetime.fromtimestamp(ts / 1000)
-        else:
-            return None
+        return _timestamp_ms_to_zoned_datetime(ts)
 
 class DSSCodeStudioTemplate(object):
     """

Diff for: dataikuapi/dss/analysis.py

@@ -260,6 +260,40 @@ def create_timeseries_forecasting_ml_task(self, target_variable,
             mltask.wait_guess_complete()
         return mltask
 
+    def create_causal_prediction_ml_task(self,
+                                         outcome_variable,
+                                         treatment_variable,
+                                         prediction_type=None,
+                                         wait_guess_complete=True):
+        """Creates a new causal prediction task in this visual analysis lab
+        for a dataset.
+
+        :param string outcome_variable: Variable to predict
+        :param string treatment_variable: Treatment variable
+        :param string prediction_type: The type of prediction problem, guessed if not provided. Valid values: CAUSAL_BINARY_CLASSIFICATION, CAUSAL_REGRESSION
+        :param boolean wait_guess_complete: if False, the returned ML task will be in 'guessing' state, i.e. analyzing the input dataset to determine feature handling and algorithms.
+                                            You should wait for the guessing to be completed by calling
+                                            ``wait_guess_complete`` on the returned object before doing anything
+                                            else (in particular calling ``train`` or ``get_settings``)
+        :return :class dataiku.dss.ml.DSSMLTask
+        """
+
+        obj = {
+            "taskType": "PREDICTION",
+            "targetVariable": outcome_variable,
+            "treatmentVariable": treatment_variable,
+            "backendType": "PY_MEMORY",
+            "guessPolicy":  "CAUSAL_PREDICTION"
+        }
+        if prediction_type is not None:
+            obj["predictionType"] = prediction_type
+        ref = self.client._perform_json("POST", "/projects/%s/lab/%s/models/" % (self.project_key, self.analysis_id), body=obj)
+        mltask = DSSMLTask(self.client, self.project_key, self.analysis_id, ref["mlTaskId"])
+
+        if wait_guess_complete:
+            mltask.wait_guess_complete()
+        return mltask
+
     def list_ml_tasks(self):
         """
         List the ML tasks in this visual analysis

Diff for: dataikuapi/dss/apideployer.py

@@ -36,23 +36,25 @@ def get_deployment(self, deployment_id):
         """
         return DSSAPIDeployerDeployment(self.client, deployment_id)
 
-    def create_deployment(self, deployment_id, service_id, infra_id, version, ignore_warnings=False):
+    def create_deployment(self, deployment_id, service_id, infra_id, version, endpoint_id=None, ignore_warnings=False):
         """
         Creates a deployment and returns the handle to interact with it. The returned deployment
         is not yet started and you need to call :meth:`~DSSAPIDeployerDeployment.start_update`
 
         :param str deployment_id: Identifier of the deployment to create
         :param str service_id: Identifier of the API Service to target
         :param str infra_id: Identifier of the deployment infrastructure to use
-        :param str version_id: Identifier of the API Service version to deploy
+        :param str version: Identifier of the API Service version to deploy
+        :param str endpoint_id: Identifier of the endpoint to deploy if you use a Deploy Anywhere infra. Ignored otherwise
         :param boolean ignore_warnings: ignore warnings concerning the governance status of the model version(s) to deploy
         :rtype: :class:`DSSAPIDeployerDeployment`
         """
         settings = {
-            "deploymentId" : deployment_id,
-            "publishedServiceId" : service_id,
-            "infraId" : infra_id,
-            "version" : version
+            "deploymentId": deployment_id,
+            "publishedServiceId": service_id,
+            "infraId": infra_id,
+            "version": version,
+            "endpointId": endpoint_id
         }
         self.client._perform_json("POST", "/api-deployer/deployments", params={"ignoreWarnings": ignore_warnings}, body=settings)
         return self.get_deployment(deployment_id)
@@ -88,7 +90,7 @@ def create_infra(self, infra_id, stage, type, govern_check_policy="NO_CHECK"):
 
         :param str infra_id: Unique Identifier of the infra to create
         :param str stage: Infrastructure stage. Stages are configurable on each API Deployer
-        :param str type: STATIC or K8S
+        :param str type: STATIC, K8S, AZURE_ML, SAGEMAKER or VERTEX_AI
         :param str govern_check_policy: PREVENT, WARN, or NO_CHECK depending if the deployer will check whether the saved model versions deployed on this infrastructure has to be managed and approved in Dataiku Govern
         :rtype: :class:`DSSAPIDeployerInfra`
         """

Diff for: dataikuapi/dss/apiservice.py

@@ -123,7 +123,6 @@ def endpoints(self):
         :rtype: list[dict]
         """
         return self.settings["endpoints"]
-    
 
     def add_prediction_endpoint(self, endpoint_id, saved_model_id):
         """
@@ -164,6 +163,22 @@ def add_forecasting_endpoint(self, endpoint_id, saved_model_id):
             "modelRef": saved_model_id
         })
 
+    def add_causal_prediction_endpoint(self, endpoint_id, saved_model_id, compute_propensity=False):
+        """
+        Add a new "visual causal prediction" endpoint to this API service.
+
+        :param string endpoint_id: identifier of the new endpoint to create
+        :param string saved_model_id: identifier of the saved model (that is currently deployed to the Flow) to use
+        :param bool compute_propensity: whether propensity should be computed, if True, the model must have a trained propensity model
+        """
+        self.settings["endpoints"].append({
+            "id": endpoint_id,
+            "type": "STD_CAUSAL_PREDICTION",
+            "modelRef": saved_model_id,
+            "computePropensity": compute_propensity,
+            "useJava": False  # Not supported for causal predictions
+        })
+
     def save(self):
         """
         Save back these settings to the API Service.

Diff for: dataikuapi/dss/codestudio.py

@@ -1,4 +1,4 @@
-from ..utils import DataikuException
+from ..utils import DataikuException, _timestamp_ms_to_zoned_datetime
 import json
 from datetime import datetime
 from .future import DSSFuture
@@ -213,6 +213,17 @@ def push_to_code_studio(self, zone):
         """
         return self.client._perform_json("GET", "/projects/%s/code-studios/%s/push/%s" % (self.project_key, self.code_studio_id, zone))
 
+    def change_owner(self, new_owner):
+        """
+        Allows to change the owner of the Code Studio
+         .. note::
+
+        only admins are allowed to change the owner of a code studio.
+
+        :param str new_owner: the id of the new owner
+        """
+        self.client._perform_json("POST", "/projects/%s/code-studios/%s/change-owner?newOwner=%s"  % (self.project_key, self.code_studio_id, new_owner))
+
 class DSSCodeStudioObjectConflicts(dict):
     """
     Summary of the conflicts on zones of a code studio.
@@ -428,10 +439,7 @@ def last_state_change(self):
         :rtype: `datetime.datetime`
         """
         ts = self.status.get("lastStateChange", 0)
-        if ts > 0:
-            return datetime.fromtimestamp(ts / 1000)
-        else:
-            return None
+        return _timestamp_ms_to_zoned_datetime(ts)
 
     def get_zones(self, as_type="names"):
         """