Repository: datarobot/airflow-provider-datarobot
Pull request #183: [PRED-11014] Adding predictions generation operators and sensor
Author: login: kessler-frost
Created: 2025-03-04T15:55:35Z | Merged: 2025-03-06T15:02:13Z

# This repository is public. Do not put here any private DataRobot or customer's data: code, datasets, model artifacts, .etc.

## Summary

Adding predictions generation airflow operators to provide similar functionality as [here](https://datarobot-public-api-client.readthedocs-hosted.com/en/latest-release/reference/predictions/predict_job.html) by the DR python SDK.



Diff:
diff --git a/datarobot_provider/operators/predictions.py b/datarobot_provider/operators/predictions.py
new file mode 100644
index 00000000..3cb63511
--- /dev/null
+++ b/datarobot_provider/operators/predictions.py
@@ -0,0 +1,120 @@
+# Copyright 2025 DataRobot, Inc. and its affiliates.
+#
+# All rights reserved.
+#
+# This is proprietary source code of DataRobot, Inc. and its affiliates.
+#
+# Released under the terms of DataRobot Tool and Utility Agreement.
+from typing import Any
+from typing import List
+from typing import Optional
+
+import datarobot as dr
+from airflow.exceptions import AirflowFailException
+from airflow.utils.context import Context
+from pandas import DataFrame
+
+from datarobot_provider.operators.base_datarobot_operator import BaseDatarobotOperator
+
+
+class RequestPredictionsOperator(BaseDatarobotOperator):
+    """
+    Starts predictions generation job for a given model and dataset.
+
+    Args:
+        project_id (str): DataRobot project ID.
+        model_id (str): DataRobot model ID.
+        dataset_id (Optional[str]): DataRobot dataset ID. Defaults to None.
+        file_path (Optional[str]): File path to the dataset. Defaults to None.
+
+    Returns:
+        str: PredictJob ID.
+
+    Raises:
+        AirflowFailException: If neither dataset_id nor file_path is provided.
+        AirflowFailException: If both dataset_id and file_path are provided.
+    """
+
+    template_fields: List[str] = ["project_id", "model_id", "dataset_id", "file_path"]
+
+    def __init__(
+        self,
+        *,
+        project_id: str,
+        model_id: str,
+        dataset_id: Optional[str] = None,
+        file_path: Optional[str] = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.model_id = model_id
+        self.dataset_id = dataset_id
+        self.file_path = file_path
+
+    def validate(self) -> None:
+        if not self.dataset_id and not self.file_path:
+            raise AirflowFailException("Either dataset_id or file_path must be provided")
+        if self.dataset_id and self.file_path:
+            raise AirflowFailException("Both dataset_id and file_path cannot be provided")
+
+    def execute(self, context: Context) -> str:
+        model = dr.Model.get(
+            project=self.project_id,
+            model_id=self.model_id,
+        )
+
+        if self.dataset_id:
+            self.log.info(
+                f"Requesting predictions for model {self.model_id} using dataset {self.dataset_id}"
+            )
+            predict_job = model.request_predictions(dataset_id=self.dataset_id)
+        else:
+            self.log.info(
+                f"Requesting predictions for model {self.model_id} using file {self.file_path}"
+            )
+            predict_job = model.request_predictions(file_path=self.file_path)
+
+        self.log.info(f"Model predictions requested, predict_job_id={predict_job.id}")
+        return predict_job.id
+
+
+class SavePredictionsToDatasetOperator(BaseDatarobotOperator):
+    """
+    Save the predictions from a PredictJob to a dataset.
+
+    Args:
+        project_id (str): DataRobot project ID.
+        predict_job_id (str): PredictJob ID.
+
+    Returns:
+        id: Dataset ID.
+
+    Raises:
+        AirflowFailException: If predict_job_id or project_id is not provided.
+    """
+
+    template_fields: List[str] = ["project_id", "predict_job_id"]
+
+    def __init__(self, *, project_id: str, predict_job_id: str, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.predict_job_id = predict_job_id
+
+    def validate(self) -> None:
+        if not self.predict_job_id or not self.project_id:
+            raise AirflowFailException("predict_job_id and project_id must be provided")
+
+    def execute(self, context: Context) -> DataFrame:
+        predictions = dr.PredictJob.get_predictions(
+            project_id=self.project_id,
+            predict_job_id=self.predict_job_id,
+        )
+
+        self.log.info(f"Predictions retrieved, predict_job_id={self.predict_job_id}")
+
+        project = dr.Project.get(self.project_id)
+        predictions_dataset = project.upload_dataset(
+            sourcedata=predictions, dataset_filename=f"prediction-results-{self.predict_job_id}.csv"
+        )
+        return predictions_dataset.id
diff --git a/datarobot_provider/sensors/predictions.py b/datarobot_provider/sensors/predictions.py
new file mode 100644
index 00000000..7f55b57d
--- /dev/null
+++ b/datarobot_provider/sensors/predictions.py
@@ -0,0 +1,56 @@
+# Copyright 2025 DataRobot, Inc. and its affiliates.
+#
+# All rights reserved.
+#
+# This is proprietary source code of DataRobot, Inc. and its affiliates.
+#
+# Released under the terms of DataRobot Tool and Utility Agreement.
+from typing import Any
+
+import datarobot as dr
+from airflow.sensors.base import BaseSensorOperator
+from airflow.utils.context import Context
+
+from datarobot_provider.constants import DATAROBOT_CONN_ID
+from datarobot_provider.hooks.datarobot import DataRobotHook
+
+
+class WaitForPredictionsSensor(BaseSensorOperator):
+    """
+    Checks whether a predictions job is complete.
+
+    Args:
+        project_id (str): DataRobot project ID.
+        predict_job_id (str): Predict job ID.
+
+    Returns:
+        bool: False if not yet completed, True when complete.
+    """
+
+    template_fields = ["project_id", "predict_job_id"]
+
+    def __init__(
+        self,
+        *,
+        project_id: str,
+        predict_job_id: str,
+        datarobot_conn_id: str = DATAROBOT_CONN_ID,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.predict_job_id = predict_job_id
+
+        self.datarobot_conn_id = datarobot_conn_id
+        self.hook = DataRobotHook(datarobot_conn_id)
+
+    def poke(self, context: Context) -> bool:
+        # Initialize DataRobot client
+        DataRobotHook(datarobot_conn_id=self.datarobot_conn_id).run()
+
+        try:
+            self.log.info("Checking if PredictJob is finished.")
+            dr.PredictJob.get(project_id=self.project_id, predict_job_id=self.predict_job_id)
+        except dr.errors.PendingJobFinished:
+            return True
+        return False
diff --git a/docs/autodoc_operators/operators_predictions.md b/docs/autodoc_operators/operators_predictions.md
new file mode 100644
index 00000000..cc28dc70
--- /dev/null
+++ b/docs/autodoc_operators/operators_predictions.md
@@ -0,0 +1,13 @@
+(predictions-modules)=
+
+# Predictions Modules
+
+```{eval-rst}
+.. autoclass:: datarobot_provider.operators.predictions.RequestPredictionsOperator
+   :members:
+```
+
+```{eval-rst}
+.. autoclass:: datarobot_provider.operators.predictions.SavePredictionsToDatasetOperator
+   :members:
+```
diff --git a/docs/autodoc_sensors/sensors_predictions.md b/docs/autodoc_sensors/sensors_predictions.md
new file mode 100644
index 00000000..a5fe8bbb
--- /dev/null
+++ b/docs/autodoc_sensors/sensors_predictions.md
@@ -0,0 +1,8 @@
+(predictions-sensors)=
+
+# Predictions Sensors
+
+```{eval-rst}
+.. autoclass:: datarobot_provider.sensors.predictions.WaitForPredictionsSensor
+   :members:
+```


Reviews:
user: login: andrius-senulis
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-04T20:00:48Z
user: login: mjnitz02
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-04T20:51:38Z
user: login: kessler-frost
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-05T10:14:46Z
user: login: kessler-frost
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-05T10:16:58Z
user: login: kessler-frost
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-05T10:21:31Z
user: login: kessler-frost
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-05T10:22:16Z
user: login: kessler-frost
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-05T10:23:17Z
user: login: kessler-frost
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-05T10:24:23Z
user: login: andrius-senulis
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-05T10:41:44Z
user: login: andrius-senulis
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-05T10:51:12Z
user: login: andrius-senulis
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-05T10:56:24Z
user: login: kessler-frost
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-05T11:01:36Z
user: login: kessler-frost
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-05T11:03:08Z
user: login: kessler-frost
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-05T11:15:06Z
user: login: kessler-frost
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-05T11:50:22Z
user: login: kessler-frost
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-05T11:51:18Z
user: login: mjnitz02
body: 👍 Outside of the get predictions operator this LGTM.  I'm not 100% sure what to do with that one related to how XCOM works.

Can you add the `autodoc` stubs in `docs/autodoc_operators` and `docs/autodoc_sensors`?  You can just copy one of the existing MD files and add the references to the predictions and they'll autopopulate.
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-03-05T14:16:44Z

Review comments:
diff_hunk: @@ -0,0 +1,202 @@
+# Copyright 2025 DataRobot, Inc. and its affiliates.
+#
+# All rights reserved.
+#
+# This is proprietary source code of DataRobot, Inc. and its affiliates.
+#
+# Released under the terms of DataRobot Tool and Utility Agreement.
+from typing import Any
+from typing import List
+from typing import Optional
+
+import datarobot as dr
+from airflow.exceptions import AirflowFailException
+from airflow.utils.context import Context
+from pandas import DataFrame
+
+from datarobot_provider.operators.base_datarobot_operator import BaseDatarobotOperator
+
+
+class GeneratePredictionsOperator(BaseDatarobotOperator):
+    """
+    Starts predictions generation job for a given model and dataset.
+
+    Args:
+        project_id (str): DataRobot project ID.
+        model_id (str): DataRobot model ID.
+        dataset_id (Optional[str]): DataRobot dataset ID. Defaults to None.
+        file_path (Optional[str]): File path to the dataset. Defaults to None.
+
+    Returns:
+        str: PredictJob ID.
+
+    Raises:
+        AirflowFailException: If neither dataset_id nor file_path is provided.
+        AirflowFailException: If both dataset_id and file_path are provided.
+    """
+
+    template_fields: List[str] = ["project_id", "model_id", "dataset_id", "file_path"]
+
+    def __init__(
+        self,
+        *,
+        project_id: str,
+        model_id: str,
+        dataset_id: Optional[str] = None,
+        file_path: Optional[str] = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.model_id = model_id
+        self.dataset_id = dataset_id
+        self.file_path = file_path
+
+    def validate(self) -> None:
+        if not self.dataset_id and not self.file_path:
+            raise AirflowFailException("Either dataset_id or file_path must be provided")
+        if self.dataset_id and self.file_path:
+            raise AirflowFailException("Both dataset_id and file_path cannot be provided")
+
+    def execute(self, context: Context) -> str:
+        model = dr.Model.get(
+            project=self.project_id,
+            model_id=self.model_id,
+        )
+
+        if self.dataset_id:
+            self.log.info(
+                f"Requesting predictions for model {self.model_id} using dataset {self.dataset_id}"
+            )
+            predict_job = model.request_predictions(dataset_id=self.dataset_id)
+        else:
+            self.log.info(
+                f"Requesting predictions for model {self.model_id} using file {self.file_path}"
+            )
+            predict_job = model.request_predictions(file_path=self.file_path)
+
+        self.log.info(f"Model predictions requested, predict_job_id={predict_job.id}")
+        return predict_job.id
+
+
+class GetPredictionsOperator(BaseDatarobotOperator):
+    """
+    Gets predictions from a PredictJob.
+
+    Args:
+        project_id (str): DataRobot project ID.
+        predict_job_id (str): PredictJob ID.
+
+    Returns:
+        DataFrame: Predictions DataFrame.
+
+    Raises:
+        AirflowFailException: If predict_job_id or project_id is not provided.
+    """
+
+    template_fields: List[str] = ["project_id", "predict_job_id"]
+
+    def __init__(self, *, project_id: str, predict_job_id: str, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.predict_job_id = predict_job_id
+
+    def validate(self) -> None:
+        if not self.predict_job_id or not self.project_id:
+            raise AirflowFailException("predict_job_id and project_id must be provided")
+
+    def execute(self, context: Context) -> DataFrame:
+        predictions = dr.PredictJob.get_predictions(
+            project_id=self.project_id,
+            predict_job_id=self.predict_job_id,
+        )
+
+        self.log.info(f"Predictions retrieved, predict_job_id={self.predict_job_id}")
+
+        ## NOTE: Can we even return a DataFrame here?
+        return predictions
+
+
+class GetPredictJobOperator(BaseDatarobotOperator):
+    """
+    Gets a PredictJob.
+
+    Args:
+        project_id (str): DataRobot project ID.
+        predict_job_id (str): PredictJob ID.
+
+    Returns:
+        dr.PredictJob: PredictJob object.
+
+    Raises:
+        AirflowFailException: If predict_job_id or project_id is not provided.
+        AirflowFailException: If PredictJob has already finished.
+    """
+
+    template_fields: List[str] = ["project_id", "predict_job_id"]
+
+    def __init__(self, *, project_id: str, predict_job_id: str, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.predict_job_id = predict_job_id
+
+    def validate(self) -> None:
+        if not self.predict_job_id or not self.project_id:
+            raise AirflowFailException("predict_job_id and project_id must be provided")
+
+    def execute(self, context: Context) -> dr.PredictJob:
+        try:
+            predict_job = dr.PredictJob.get(
+                project_id=self.project_id,
+                predict_job_id=self.predict_job_id,
+            )
+        except dr.errors.PendingJobFinished:
+            # NOTE: Is this the correct exception to raise in this case?
+            raise AirflowFailException(
+                "PredictJob has already finished, please use PredictJobGetPredictionsOperator to retrieve predictions"
+            ) from dr.errors.PendingJobFinished
+
+        self.log.info(f"PredictJob retrieved, predict_job_id={self.predict_job_id}")
+
+        # NOTE: Can we return the PredictJob object here?
+        return predict_job
+
+
+class AwaitAndGetPredictionsOperator(BaseDatarobotOperator):
+    """
+    Waits and retrieves the predictions from a PredictJob.
+
+    Args:
+        project_id (str): DataRobot project ID.
+        predict_job_id (str): PredictJob ID.
+
+    Returns:
+        DataFrame: Predictions DataFrame.
+
+    Raises:
+        AirflowFailException: If predict_job_id or project_id is not provided.
+    """
+
+    template_fields: List[str] = ["project_id", "predict_job_id"]
+
+    def __init__(self, *, project_id: str, predict_job_id: str, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.predict_job_id = predict_job_id
+
+    def validate(self) -> None:
+        if not self.predict_job_id or not self.project_id:
+            raise AirflowFailException("predict_job_id and project_id must be provided")
+
+    def execute(self, context: Context) -> DataFrame:
+        predict_job = dr.PredictJob.get(
+            project_id=self.project_id,
+            predict_job_id=self.predict_job_id,
+        )
+
+        predictions = predict_job.get_result_when_complete()
+
+        self.log.info(f"Predictions waited for and retrieved, predict_job_id={self.predict_job_id}")
+
+        # NOTE: Again, can we return a DataFrame here?
+        return predictions
path: datarobot_provider/operators/predictions.py
user: login: mjnitz02
body: It might be better to just return the job id here. I'm not exactly sure if airflow can handle a dataframe and what it would do with the resulting dataframe
created_at: 2025-03-04T19:02:37Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,202 @@
+# Copyright 2025 DataRobot, Inc. and its affiliates.
+#
+# All rights reserved.
+#
+# This is proprietary source code of DataRobot, Inc. and its affiliates.
+#
+# Released under the terms of DataRobot Tool and Utility Agreement.
+from typing import Any
+from typing import List
+from typing import Optional
+
+import datarobot as dr
+from airflow.exceptions import AirflowFailException
+from airflow.utils.context import Context
+from pandas import DataFrame
+
+from datarobot_provider.operators.base_datarobot_operator import BaseDatarobotOperator
+
+
+class GeneratePredictionsOperator(BaseDatarobotOperator):
+    """
+    Starts predictions generation job for a given model and dataset.
+
+    Args:
+        project_id (str): DataRobot project ID.
+        model_id (str): DataRobot model ID.
+        dataset_id (Optional[str]): DataRobot dataset ID. Defaults to None.
+        file_path (Optional[str]): File path to the dataset. Defaults to None.
+
+    Returns:
+        str: PredictJob ID.
+
+    Raises:
+        AirflowFailException: If neither dataset_id nor file_path is provided.
+        AirflowFailException: If both dataset_id and file_path are provided.
+    """
+
+    template_fields: List[str] = ["project_id", "model_id", "dataset_id", "file_path"]
+
+    def __init__(
+        self,
+        *,
+        project_id: str,
+        model_id: str,
+        dataset_id: Optional[str] = None,
+        file_path: Optional[str] = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.model_id = model_id
+        self.dataset_id = dataset_id
+        self.file_path = file_path
+
+    def validate(self) -> None:
+        if not self.dataset_id and not self.file_path:
+            raise AirflowFailException("Either dataset_id or file_path must be provided")
+        if self.dataset_id and self.file_path:
+            raise AirflowFailException("Both dataset_id and file_path cannot be provided")
+
+    def execute(self, context: Context) -> str:
+        model = dr.Model.get(
+            project=self.project_id,
+            model_id=self.model_id,
+        )
+
+        if self.dataset_id:
+            self.log.info(
+                f"Requesting predictions for model {self.model_id} using dataset {self.dataset_id}"
+            )
+            predict_job = model.request_predictions(dataset_id=self.dataset_id)
+        else:
+            self.log.info(
+                f"Requesting predictions for model {self.model_id} using file {self.file_path}"
+            )
+            predict_job = model.request_predictions(file_path=self.file_path)
+
+        self.log.info(f"Model predictions requested, predict_job_id={predict_job.id}")
+        return predict_job.id
+
+
+class GetPredictionsOperator(BaseDatarobotOperator):
+    """
+    Gets predictions from a PredictJob.
+
+    Args:
+        project_id (str): DataRobot project ID.
+        predict_job_id (str): PredictJob ID.
+
+    Returns:
+        DataFrame: Predictions DataFrame.
+
+    Raises:
+        AirflowFailException: If predict_job_id or project_id is not provided.
+    """
+
+    template_fields: List[str] = ["project_id", "predict_job_id"]
+
+    def __init__(self, *, project_id: str, predict_job_id: str, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.predict_job_id = predict_job_id
+
+    def validate(self) -> None:
+        if not self.predict_job_id or not self.project_id:
+            raise AirflowFailException("predict_job_id and project_id must be provided")
+
+    def execute(self, context: Context) -> DataFrame:
+        predictions = dr.PredictJob.get_predictions(
+            project_id=self.project_id,
+            predict_job_id=self.predict_job_id,
+        )
+
+        self.log.info(f"Predictions retrieved, predict_job_id={self.predict_job_id}")
+
+        ## NOTE: Can we even return a DataFrame here?
+        return predictions
+
+
+class GetPredictJobOperator(BaseDatarobotOperator):
+    """
+    Gets a PredictJob.
+
+    Args:
+        project_id (str): DataRobot project ID.
+        predict_job_id (str): PredictJob ID.
+
+    Returns:
+        dr.PredictJob: PredictJob object.
+
+    Raises:
+        AirflowFailException: If predict_job_id or project_id is not provided.
+        AirflowFailException: If PredictJob has already finished.
+    """
+
+    template_fields: List[str] = ["project_id", "predict_job_id"]
+
+    def __init__(self, *, project_id: str, predict_job_id: str, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.predict_job_id = predict_job_id
+
+    def validate(self) -> None:
+        if not self.predict_job_id or not self.project_id:
+            raise AirflowFailException("predict_job_id and project_id must be provided")
+
+    def execute(self, context: Context) -> dr.PredictJob:
+        try:
+            predict_job = dr.PredictJob.get(
+                project_id=self.project_id,
+                predict_job_id=self.predict_job_id,
+            )
+        except dr.errors.PendingJobFinished:
+            # NOTE: Is this the correct exception to raise in this case?
+            raise AirflowFailException(
+                "PredictJob has already finished, please use PredictJobGetPredictionsOperator to retrieve predictions"
+            ) from dr.errors.PendingJobFinished
+
+        self.log.info(f"PredictJob retrieved, predict_job_id={self.predict_job_id}")
+
+        # NOTE: Can we return the PredictJob object here?
+        return predict_job
path: datarobot_provider/operators/predictions.py
user: login: mjnitz02
body: We can only return the id of it, the `XCOM` can't handle the object
created_at: 2025-03-04T19:02:56Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,202 @@
+# Copyright 2025 DataRobot, Inc. and its affiliates.
+#
+# All rights reserved.
+#
+# This is proprietary source code of DataRobot, Inc. and its affiliates.
+#
+# Released under the terms of DataRobot Tool and Utility Agreement.
+from typing import Any
+from typing import List
+from typing import Optional
+
+import datarobot as dr
+from airflow.exceptions import AirflowFailException
+from airflow.utils.context import Context
+from pandas import DataFrame
+
+from datarobot_provider.operators.base_datarobot_operator import BaseDatarobotOperator
+
+
+class GeneratePredictionsOperator(BaseDatarobotOperator):
+    """
+    Starts predictions generation job for a given model and dataset.
+
+    Args:
+        project_id (str): DataRobot project ID.
+        model_id (str): DataRobot model ID.
+        dataset_id (Optional[str]): DataRobot dataset ID. Defaults to None.
+        file_path (Optional[str]): File path to the dataset. Defaults to None.
+
+    Returns:
+        str: PredictJob ID.
+
+    Raises:
+        AirflowFailException: If neither dataset_id nor file_path is provided.
+        AirflowFailException: If both dataset_id and file_path are provided.
+    """
+
+    template_fields: List[str] = ["project_id", "model_id", "dataset_id", "file_path"]
+
+    def __init__(
+        self,
+        *,
+        project_id: str,
+        model_id: str,
+        dataset_id: Optional[str] = None,
+        file_path: Optional[str] = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.model_id = model_id
+        self.dataset_id = dataset_id
+        self.file_path = file_path
+
+    def validate(self) -> None:
+        if not self.dataset_id and not self.file_path:
+            raise AirflowFailException("Either dataset_id or file_path must be provided")
+        if self.dataset_id and self.file_path:
+            raise AirflowFailException("Both dataset_id and file_path cannot be provided")
+
+    def execute(self, context: Context) -> str:
+        model = dr.Model.get(
+            project=self.project_id,
+            model_id=self.model_id,
+        )
+
+        if self.dataset_id:
+            self.log.info(
+                f"Requesting predictions for model {self.model_id} using dataset {self.dataset_id}"
+            )
+            predict_job = model.request_predictions(dataset_id=self.dataset_id)
+        else:
+            self.log.info(
+                f"Requesting predictions for model {self.model_id} using file {self.file_path}"
+            )
+            predict_job = model.request_predictions(file_path=self.file_path)
+
+        self.log.info(f"Model predictions requested, predict_job_id={predict_job.id}")
+        return predict_job.id
+
+
+class GetPredictionsOperator(BaseDatarobotOperator):
+    """
+    Gets predictions from a PredictJob.
+
+    Args:
+        project_id (str): DataRobot project ID.
+        predict_job_id (str): PredictJob ID.
+
+    Returns:
+        DataFrame: Predictions DataFrame.
+
+    Raises:
+        AirflowFailException: If predict_job_id or project_id is not provided.
+    """
+
+    template_fields: List[str] = ["project_id", "predict_job_id"]
+
+    def __init__(self, *, project_id: str, predict_job_id: str, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.predict_job_id = predict_job_id
+
+    def validate(self) -> None:
+        if not self.predict_job_id or not self.project_id:
+            raise AirflowFailException("predict_job_id and project_id must be provided")
+
+    def execute(self, context: Context) -> DataFrame:
+        predictions = dr.PredictJob.get_predictions(
+            project_id=self.project_id,
+            predict_job_id=self.predict_job_id,
+        )
+
+        self.log.info(f"Predictions retrieved, predict_job_id={self.predict_job_id}")
+
+        ## NOTE: Can we even return a DataFrame here?
+        return predictions
path: datarobot_provider/operators/predictions.py
user: login: andrius-senulis
body: [Airflow docs](https://airflow.apache.org/docs/apache-airflow/stable/core-concepts/xcoms.html#xcoms) suggests that large data should not be passed around in XComs. Basically, operators should return small serializable data: strings, ints, dicts, etc.

This operator should persist the data somewhere, most likely in S3 which should be configurable in the operator. I wonder how is this operator will be used, do you know?
created_at: 2025-03-04T19:48:50Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,202 @@
+# Copyright 2025 DataRobot, Inc. and its affiliates.
+#
+# All rights reserved.
+#
+# This is proprietary source code of DataRobot, Inc. and its affiliates.
+#
+# Released under the terms of DataRobot Tool and Utility Agreement.
+from typing import Any
+from typing import List
+from typing import Optional
+
+import datarobot as dr
+from airflow.exceptions import AirflowFailException
+from airflow.utils.context import Context
+from pandas import DataFrame
+
+from datarobot_provider.operators.base_datarobot_operator import BaseDatarobotOperator
+
+
+class GeneratePredictionsOperator(BaseDatarobotOperator):
path: datarobot_provider/operators/predictions.py
user: login: andrius-senulis
body: For context, we call this way of making predictions - leaderboard predictions. It's making predictions directly against the model without deploying it.

I would prefer to call it something else than `GeneratePredictions`. `RequestPredictions` makes most sense imho, which also matches the call `model.request_predictions()`.
created_at: 2025-03-04T19:48:52Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,202 @@
+# Copyright 2025 DataRobot, Inc. and its affiliates.
+#
+# All rights reserved.
+#
+# This is proprietary source code of DataRobot, Inc. and its affiliates.
+#
+# Released under the terms of DataRobot Tool and Utility Agreement.
+from typing import Any
+from typing import List
+from typing import Optional
+
+import datarobot as dr
+from airflow.exceptions import AirflowFailException
+from airflow.utils.context import Context
+from pandas import DataFrame
+
+from datarobot_provider.operators.base_datarobot_operator import BaseDatarobotOperator
+
+
+class GeneratePredictionsOperator(BaseDatarobotOperator):
+    """
+    Starts predictions generation job for a given model and dataset.
+
+    Args:
+        project_id (str): DataRobot project ID.
+        model_id (str): DataRobot model ID.
+        dataset_id (Optional[str]): DataRobot dataset ID. Defaults to None.
+        file_path (Optional[str]): File path to the dataset. Defaults to None.
+
+    Returns:
+        str: PredictJob ID.
+
+    Raises:
+        AirflowFailException: If neither dataset_id nor file_path is provided.
+        AirflowFailException: If both dataset_id and file_path are provided.
+    """
+
+    template_fields: List[str] = ["project_id", "model_id", "dataset_id", "file_path"]
+
+    def __init__(
+        self,
+        *,
+        project_id: str,
+        model_id: str,
+        dataset_id: Optional[str] = None,
+        file_path: Optional[str] = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.model_id = model_id
+        self.dataset_id = dataset_id
+        self.file_path = file_path
+
+    def validate(self) -> None:
+        if not self.dataset_id and not self.file_path:
+            raise AirflowFailException("Either dataset_id or file_path must be provided")
+        if self.dataset_id and self.file_path:
+            raise AirflowFailException("Both dataset_id and file_path cannot be provided")
+
+    def execute(self, context: Context) -> str:
+        model = dr.Model.get(
+            project=self.project_id,
+            model_id=self.model_id,
+        )
+
+        if self.dataset_id:
+            self.log.info(
+                f"Requesting predictions for model {self.model_id} using dataset {self.dataset_id}"
+            )
+            predict_job = model.request_predictions(dataset_id=self.dataset_id)
+        else:
+            self.log.info(
+                f"Requesting predictions for model {self.model_id} using file {self.file_path}"
+            )
+            predict_job = model.request_predictions(file_path=self.file_path)
+
+        self.log.info(f"Model predictions requested, predict_job_id={predict_job.id}")
+        return predict_job.id
+
+
+class GetPredictionsOperator(BaseDatarobotOperator):
+    """
+    Gets predictions from a PredictJob.
+
+    Args:
+        project_id (str): DataRobot project ID.
+        predict_job_id (str): PredictJob ID.
+
+    Returns:
+        DataFrame: Predictions DataFrame.
+
+    Raises:
+        AirflowFailException: If predict_job_id or project_id is not provided.
+    """
+
+    template_fields: List[str] = ["project_id", "predict_job_id"]
+
+    def __init__(self, *, project_id: str, predict_job_id: str, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.predict_job_id = predict_job_id
+
+    def validate(self) -> None:
+        if not self.predict_job_id or not self.project_id:
+            raise AirflowFailException("predict_job_id and project_id must be provided")
+
+    def execute(self, context: Context) -> DataFrame:
+        predictions = dr.PredictJob.get_predictions(
+            project_id=self.project_id,
+            predict_job_id=self.predict_job_id,
+        )
+
+        self.log.info(f"Predictions retrieved, predict_job_id={self.predict_job_id}")
+
+        ## NOTE: Can we even return a DataFrame here?
+        return predictions
+
+
+class GetPredictJobOperator(BaseDatarobotOperator):
+    """
+    Gets a PredictJob.
+
+    Args:
+        project_id (str): DataRobot project ID.
+        predict_job_id (str): PredictJob ID.
+
+    Returns:
+        dr.PredictJob: PredictJob object.
+
+    Raises:
+        AirflowFailException: If predict_job_id or project_id is not provided.
+        AirflowFailException: If PredictJob has already finished.
+    """
+
+    template_fields: List[str] = ["project_id", "predict_job_id"]
+
+    def __init__(self, *, project_id: str, predict_job_id: str, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.predict_job_id = predict_job_id
+
+    def validate(self) -> None:
+        if not self.predict_job_id or not self.project_id:
+            raise AirflowFailException("predict_job_id and project_id must be provided")
+
+    def execute(self, context: Context) -> dr.PredictJob:
+        try:
+            predict_job = dr.PredictJob.get(
+                project_id=self.project_id,
+                predict_job_id=self.predict_job_id,
+            )
+        except dr.errors.PendingJobFinished:
+            # NOTE: Is this the correct exception to raise in this case?
+            raise AirflowFailException(
+                "PredictJob has already finished, please use PredictJobGetPredictionsOperator to retrieve predictions"
+            ) from dr.errors.PendingJobFinished
+
+        self.log.info(f"PredictJob retrieved, predict_job_id={self.predict_job_id}")
+
+        # NOTE: Can we return the PredictJob object here?
+        return predict_job
path: datarobot_provider/operators/predictions.py
user: login: andrius-senulis
body: No, it should be job ID.
created_at: 2025-03-04T19:53:31Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,202 @@
+# Copyright 2025 DataRobot, Inc. and its affiliates.
+#
+# All rights reserved.
+#
+# This is proprietary source code of DataRobot, Inc. and its affiliates.
+#
+# Released under the terms of DataRobot Tool and Utility Agreement.
+from typing import Any
+from typing import List
+from typing import Optional
+
+import datarobot as dr
+from airflow.exceptions import AirflowFailException
+from airflow.utils.context import Context
+from pandas import DataFrame
+
+from datarobot_provider.operators.base_datarobot_operator import BaseDatarobotOperator
+
+
+class GeneratePredictionsOperator(BaseDatarobotOperator):
+    """
+    Starts predictions generation job for a given model and dataset.
+
+    Args:
+        project_id (str): DataRobot project ID.
+        model_id (str): DataRobot model ID.
+        dataset_id (Optional[str]): DataRobot dataset ID. Defaults to None.
+        file_path (Optional[str]): File path to the dataset. Defaults to None.
+
+    Returns:
+        str: PredictJob ID.
+
+    Raises:
+        AirflowFailException: If neither dataset_id nor file_path is provided.
+        AirflowFailException: If both dataset_id and file_path are provided.
+    """
+
+    template_fields: List[str] = ["project_id", "model_id", "dataset_id", "file_path"]
+
+    def __init__(
+        self,
+        *,
+        project_id: str,
+        model_id: str,
+        dataset_id: Optional[str] = None,
+        file_path: Optional[str] = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.model_id = model_id
+        self.dataset_id = dataset_id
+        self.file_path = file_path
+
+    def validate(self) -> None:
+        if not self.dataset_id and not self.file_path:
+            raise AirflowFailException("Either dataset_id or file_path must be provided")
+        if self.dataset_id and self.file_path:
+            raise AirflowFailException("Both dataset_id and file_path cannot be provided")
+
+    def execute(self, context: Context) -> str:
+        model = dr.Model.get(
+            project=self.project_id,
+            model_id=self.model_id,
+        )
+
+        if self.dataset_id:
+            self.log.info(
+                f"Requesting predictions for model {self.model_id} using dataset {self.dataset_id}"
+            )
+            predict_job = model.request_predictions(dataset_id=self.dataset_id)
+        else:
+            self.log.info(
+                f"Requesting predictions for model {self.model_id} using file {self.file_path}"
+            )
+            predict_job = model.request_predictions(file_path=self.file_path)
+
+        self.log.info(f"Model predictions requested, predict_job_id={predict_job.id}")
+        return predict_job.id
+
+
+class GetPredictionsOperator(BaseDatarobotOperator):
+    """
+    Gets predictions from a PredictJob.
+
+    Args:
+        project_id (str): DataRobot project ID.
+        predict_job_id (str): PredictJob ID.
+
+    Returns:
+        DataFrame: Predictions DataFrame.
+
+    Raises:
+        AirflowFailException: If predict_job_id or project_id is not provided.
+    """
+
+    template_fields: List[str] = ["project_id", "predict_job_id"]
+
+    def __init__(self, *, project_id: str, predict_job_id: str, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.predict_job_id = predict_job_id
+
+    def validate(self) -> None:
+        if not self.predict_job_id or not self.project_id:
+            raise AirflowFailException("predict_job_id and project_id must be provided")
+
+    def execute(self, context: Context) -> DataFrame:
+        predictions = dr.PredictJob.get_predictions(
+            project_id=self.project_id,
+            predict_job_id=self.predict_job_id,
+        )
+
+        self.log.info(f"Predictions retrieved, predict_job_id={self.predict_job_id}")
+
+        ## NOTE: Can we even return a DataFrame here?
+        return predictions
+
+
+class GetPredictJobOperator(BaseDatarobotOperator):
+    """
+    Gets a PredictJob.
+
+    Args:
+        project_id (str): DataRobot project ID.
+        predict_job_id (str): PredictJob ID.
+
+    Returns:
+        dr.PredictJob: PredictJob object.
+
+    Raises:
+        AirflowFailException: If predict_job_id or project_id is not provided.
+        AirflowFailException: If PredictJob has already finished.
+    """
+
+    template_fields: List[str] = ["project_id", "predict_job_id"]
+
+    def __init__(self, *, project_id: str, predict_job_id: str, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.predict_job_id = predict_job_id
+
+    def validate(self) -> None:
+        if not self.predict_job_id or not self.project_id:
+            raise AirflowFailException("predict_job_id and project_id must be provided")
+
+    def execute(self, context: Context) -> dr.PredictJob:
+        try:
+            predict_job = dr.PredictJob.get(
+                project_id=self.project_id,
+                predict_job_id=self.predict_job_id,
+            )
+        except dr.errors.PendingJobFinished:
path: datarobot_provider/operators/predictions.py
user: login: andrius-senulis
body: In what case is this error raised?
created_at: 2025-03-04T20:00:35Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,202 @@
+# Copyright 2025 DataRobot, Inc. and its affiliates.
+#
+# All rights reserved.
+#
+# This is proprietary source code of DataRobot, Inc. and its affiliates.
+#
+# Released under the terms of DataRobot Tool and Utility Agreement.
+from typing import Any
+from typing import List
+from typing import Optional
+
+import datarobot as dr
+from airflow.exceptions import AirflowFailException
+from airflow.utils.context import Context
+from pandas import DataFrame
+
+from datarobot_provider.operators.base_datarobot_operator import BaseDatarobotOperator
+
+
+class GeneratePredictionsOperator(BaseDatarobotOperator):
+    """
+    Starts predictions generation job for a given model and dataset.
+
+    Args:
+        project_id (str): DataRobot project ID.
+        model_id (str): DataRobot model ID.
+        dataset_id (Optional[str]): DataRobot dataset ID. Defaults to None.
+        file_path (Optional[str]): File path to the dataset. Defaults to None.
+
+    Returns:
+        str: PredictJob ID.
+
+    Raises:
+        AirflowFailException: If neither dataset_id nor file_path is provided.
+        AirflowFailException: If both dataset_id and file_path are provided.
+    """
+
+    template_fields: List[str] = ["project_id", "model_id", "dataset_id", "file_path"]
+
+    def __init__(
+        self,
+        *,
+        project_id: str,
+        model_id: str,
+        dataset_id: Optional[str] = None,
+        file_path: Optional[str] = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.model_id = model_id
+        self.dataset_id = dataset_id
+        self.file_path = file_path
+
+    def validate(self) -> None:
+        if not self.dataset_id and not self.file_path:
+            raise AirflowFailException("Either dataset_id or file_path must be provided")
+        if self.dataset_id and self.file_path:
+            raise AirflowFailException("Both dataset_id and file_path cannot be provided")
+
+    def execute(self, context: Context) -> str:
+        model = dr.Model.get(
+            project=self.project_id,
+            model_id=self.model_id,
+        )
+
+        if self.dataset_id:
+            self.log.info(
+                f"Requesting predictions for model {self.model_id} using dataset {self.dataset_id}"
+            )
+            predict_job = model.request_predictions(dataset_id=self.dataset_id)
+        else:
+            self.log.info(
+                f"Requesting predictions for model {self.model_id} using file {self.file_path}"
+            )
+            predict_job = model.request_predictions(file_path=self.file_path)
+
+        self.log.info(f"Model predictions requested, predict_job_id={predict_job.id}")
+        return predict_job.id
+
+
+class GetPredictionsOperator(BaseDatarobotOperator):
+    """
+    Gets predictions from a PredictJob.
+
+    Args:
+        project_id (str): DataRobot project ID.
+        predict_job_id (str): PredictJob ID.
+
+    Returns:
+        DataFrame: Predictions DataFrame.
+
+    Raises:
+        AirflowFailException: If predict_job_id or project_id is not provided.
+    """
+
+    template_fields: List[str] = ["project_id", "predict_job_id"]
+
+    def __init__(self, *, project_id: str, predict_job_id: str, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.predict_job_id = predict_job_id
+
+    def validate(self) -> None:
+        if not self.predict_job_id or not self.project_id:
+            raise AirflowFailException("predict_job_id and project_id must be provided")
+
+    def execute(self, context: Context) -> DataFrame:
+        predictions = dr.PredictJob.get_predictions(
+            project_id=self.project_id,
+            predict_job_id=self.predict_job_id,
+        )
+
+        self.log.info(f"Predictions retrieved, predict_job_id={self.predict_job_id}")
+
+        ## NOTE: Can we even return a DataFrame here?
+        return predictions
path: datarobot_provider/operators/predictions.py
user: login: kessler-frost
body: I'm not exactly sure how will this operator be used right now, but do you think it makes sense instead of s3 - for this particular case - we can instead save it as a dataset and return the id of that dataset in the project itself? this is just a dataframe right?
created_at: 2025-03-05T10:14:46Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,202 @@
+# Copyright 2025 DataRobot, Inc. and its affiliates.
+#
+# All rights reserved.
+#
+# This is proprietary source code of DataRobot, Inc. and its affiliates.
+#
+# Released under the terms of DataRobot Tool and Utility Agreement.
+from typing import Any
+from typing import List
+from typing import Optional
+
+import datarobot as dr
+from airflow.exceptions import AirflowFailException
+from airflow.utils.context import Context
+from pandas import DataFrame
+
+from datarobot_provider.operators.base_datarobot_operator import BaseDatarobotOperator
+
+
+class GeneratePredictionsOperator(BaseDatarobotOperator):
path: datarobot_provider/operators/predictions.py
user: login: kessler-frost
body: Agreed. Yeah I was going back and forth on this one, earlier (if you see the commit history) I had planned to keep it as close to our python interface as possible, but then I saw how other operators were named and made changes accordingly. I'll update the name to `RequestPredictions` 👍 .
created_at: 2025-03-05T10:16:58Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,202 @@
+# Copyright 2025 DataRobot, Inc. and its affiliates.
+#
+# All rights reserved.
+#
+# This is proprietary source code of DataRobot, Inc. and its affiliates.
+#
+# Released under the terms of DataRobot Tool and Utility Agreement.
+from typing import Any
+from typing import List
+from typing import Optional
+
+import datarobot as dr
+from airflow.exceptions import AirflowFailException
+from airflow.utils.context import Context
+from pandas import DataFrame
+
+from datarobot_provider.operators.base_datarobot_operator import BaseDatarobotOperator
+
+
+class GeneratePredictionsOperator(BaseDatarobotOperator):
+    """
+    Starts predictions generation job for a given model and dataset.
+
+    Args:
+        project_id (str): DataRobot project ID.
+        model_id (str): DataRobot model ID.
+        dataset_id (Optional[str]): DataRobot dataset ID. Defaults to None.
+        file_path (Optional[str]): File path to the dataset. Defaults to None.
+
+    Returns:
+        str: PredictJob ID.
+
+    Raises:
+        AirflowFailException: If neither dataset_id nor file_path is provided.
+        AirflowFailException: If both dataset_id and file_path are provided.
+    """
+
+    template_fields: List[str] = ["project_id", "model_id", "dataset_id", "file_path"]
+
+    def __init__(
+        self,
+        *,
+        project_id: str,
+        model_id: str,
+        dataset_id: Optional[str] = None,
+        file_path: Optional[str] = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.model_id = model_id
+        self.dataset_id = dataset_id
+        self.file_path = file_path
+
+    def validate(self) -> None:
+        if not self.dataset_id and not self.file_path:
+            raise AirflowFailException("Either dataset_id or file_path must be provided")
+        if self.dataset_id and self.file_path:
+            raise AirflowFailException("Both dataset_id and file_path cannot be provided")
+
+    def execute(self, context: Context) -> str:
+        model = dr.Model.get(
+            project=self.project_id,
+            model_id=self.model_id,
+        )
+
+        if self.dataset_id:
+            self.log.info(
+                f"Requesting predictions for model {self.model_id} using dataset {self.dataset_id}"
+            )
+            predict_job = model.request_predictions(dataset_id=self.dataset_id)
+        else:
+            self.log.info(
+                f"Requesting predictions for model {self.model_id} using file {self.file_path}"
+            )
+            predict_job = model.request_predictions(file_path=self.file_path)
+
+        self.log.info(f"Model predictions requested, predict_job_id={predict_job.id}")
+        return predict_job.id
+
+
+class GetPredictionsOperator(BaseDatarobotOperator):
+    """
+    Gets predictions from a PredictJob.
+
+    Args:
+        project_id (str): DataRobot project ID.
+        predict_job_id (str): PredictJob ID.
+
+    Returns:
+        DataFrame: Predictions DataFrame.
+
+    Raises:
+        AirflowFailException: If predict_job_id or project_id is not provided.
+    """
+
+    template_fields: List[str] = ["project_id", "predict_job_id"]
+
+    def __init__(self, *, project_id: str, predict_job_id: str, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.predict_job_id = predict_job_id
+
+    def validate(self) -> None:
+        if not self.predict_job_id or not self.project_id:
+            raise AirflowFailException("predict_job_id and project_id must be provided")
+
+    def execute(self, context: Context) -> DataFrame:
+        predictions = dr.PredictJob.get_predictions(
+            project_id=self.project_id,
+            predict_job_id=self.predict_job_id,
+        )
+
+        self.log.info(f"Predictions retrieved, predict_job_id={self.predict_job_id}")
+
+        ## NOTE: Can we even return a DataFrame here?
+        return predictions
+
+
+class GetPredictJobOperator(BaseDatarobotOperator):
+    """
+    Gets a PredictJob.
+
+    Args:
+        project_id (str): DataRobot project ID.
+        predict_job_id (str): PredictJob ID.
+
+    Returns:
+        dr.PredictJob: PredictJob object.
+
+    Raises:
+        AirflowFailException: If predict_job_id or project_id is not provided.
+        AirflowFailException: If PredictJob has already finished.
+    """
+
+    template_fields: List[str] = ["project_id", "predict_job_id"]
+
+    def __init__(self, *, project_id: str, predict_job_id: str, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.predict_job_id = predict_job_id
+
+    def validate(self) -> None:
+        if not self.predict_job_id or not self.project_id:
+            raise AirflowFailException("predict_job_id and project_id must be provided")
+
+    def execute(self, context: Context) -> dr.PredictJob:
+        try:
+            predict_job = dr.PredictJob.get(
+                project_id=self.project_id,
+                predict_job_id=self.predict_job_id,
+            )
+        except dr.errors.PendingJobFinished:
path: datarobot_provider/operators/predictions.py
user: login: kessler-frost
body: In case the PredictJob has already finished, we don't want Airflow to keep trying to run this if the job has already finished right? And instead ask to run the relevant operator to get the predictions.

That said, we can also instead return the predictions themselves if the PredictJob has finished, but idk if that's a good way to do it since the output type will be different in different cases.
created_at: 2025-03-05T10:21:31Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,202 @@
+# Copyright 2025 DataRobot, Inc. and its affiliates.
+#
+# All rights reserved.
+#
+# This is proprietary source code of DataRobot, Inc. and its affiliates.
+#
+# Released under the terms of DataRobot Tool and Utility Agreement.
+from typing import Any
+from typing import List
+from typing import Optional
+
+import datarobot as dr
+from airflow.exceptions import AirflowFailException
+from airflow.utils.context import Context
+from pandas import DataFrame
+
+from datarobot_provider.operators.base_datarobot_operator import BaseDatarobotOperator
+
+
+class GeneratePredictionsOperator(BaseDatarobotOperator):
+    """
+    Starts predictions generation job for a given model and dataset.
+
+    Args:
+        project_id (str): DataRobot project ID.
+        model_id (str): DataRobot model ID.
+        dataset_id (Optional[str]): DataRobot dataset ID. Defaults to None.
+        file_path (Optional[str]): File path to the dataset. Defaults to None.
+
+    Returns:
+        str: PredictJob ID.
+
+    Raises:
+        AirflowFailException: If neither dataset_id nor file_path is provided.
+        AirflowFailException: If both dataset_id and file_path are provided.
+    """
+
+    template_fields: List[str] = ["project_id", "model_id", "dataset_id", "file_path"]
+
+    def __init__(
+        self,
+        *,
+        project_id: str,
+        model_id: str,
+        dataset_id: Optional[str] = None,
+        file_path: Optional[str] = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.model_id = model_id
+        self.dataset_id = dataset_id
+        self.file_path = file_path
+
+    def validate(self) -> None:
+        if not self.dataset_id and not self.file_path:
+            raise AirflowFailException("Either dataset_id or file_path must be provided")
+        if self.dataset_id and self.file_path:
+            raise AirflowFailException("Both dataset_id and file_path cannot be provided")
+
+    def execute(self, context: Context) -> str:
+        model = dr.Model.get(
+            project=self.project_id,
+            model_id=self.model_id,
+        )
+
+        if self.dataset_id:
+            self.log.info(
+                f"Requesting predictions for model {self.model_id} using dataset {self.dataset_id}"
+            )
+            predict_job = model.request_predictions(dataset_id=self.dataset_id)
+        else:
+            self.log.info(
+                f"Requesting predictions for model {self.model_id} using file {self.file_path}"
+            )
+            predict_job = model.request_predictions(file_path=self.file_path)
+
+        self.log.info(f"Model predictions requested, predict_job_id={predict_job.id}")
+        return predict_job.id
+
+
+class GetPredictionsOperator(BaseDatarobotOperator):
+    """
+    Gets predictions from a PredictJob.
+
+    Args:
+        project_id (str): DataRobot project ID.
+        predict_job_id (str): PredictJob ID.
+
+    Returns:
+        DataFrame: Predictions DataFrame.
+
+    Raises:
+        AirflowFailException: If predict_job_id or project_id is not provided.
+    """
+
+    template_fields: List[str] = ["project_id", "predict_job_id"]
+
+    def __init__(self, *, project_id: str, predict_job_id: str, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.predict_job_id = predict_job_id
+
+    def validate(self) -> None:
+        if not self.predict_job_id or not self.project_id:
+            raise AirflowFailException("predict_job_id and project_id must be provided")
+
+    def execute(self, context: Context) -> DataFrame:
+        predictions = dr.PredictJob.get_predictions(
+            project_id=self.project_id,
+            predict_job_id=self.predict_job_id,
+        )
+
+        self.log.info(f"Predictions retrieved, predict_job_id={self.predict_job_id}")
+
+        ## NOTE: Can we even return a DataFrame here?
+        return predictions
+
+
+class GetPredictJobOperator(BaseDatarobotOperator):
+    """
+    Gets a PredictJob.
+
+    Args:
+        project_id (str): DataRobot project ID.
+        predict_job_id (str): PredictJob ID.
+
+    Returns:
+        dr.PredictJob: PredictJob object.
+
+    Raises:
+        AirflowFailException: If predict_job_id or project_id is not provided.
+        AirflowFailException: If PredictJob has already finished.
+    """
+
+    template_fields: List[str] = ["project_id", "predict_job_id"]
+
+    def __init__(self, *, project_id: str, predict_job_id: str, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.predict_job_id = predict_job_id
+
+    def validate(self) -> None:
+        if not self.predict_job_id or not self.project_id:
+            raise AirflowFailException("predict_job_id and project_id must be provided")
+
+    def execute(self, context: Context) -> dr.PredictJob:
+        try:
+            predict_job = dr.PredictJob.get(
+                project_id=self.project_id,
+                predict_job_id=self.predict_job_id,
+            )
+        except dr.errors.PendingJobFinished:
+            # NOTE: Is this the correct exception to raise in this case?
+            raise AirflowFailException(
+                "PredictJob has already finished, please use PredictJobGetPredictionsOperator to retrieve predictions"
+            ) from dr.errors.PendingJobFinished
+
+        self.log.info(f"PredictJob retrieved, predict_job_id={self.predict_job_id}")
+
+        # NOTE: Can we return the PredictJob object here?
+        return predict_job
path: datarobot_provider/operators/predictions.py
user: login: kessler-frost
body: So, we are inputting the PredictJob id, and returning the same PredictJob id then?

Seems a bit weird to return the same thing the user input, without doing much. Is the point that in case it finished running then we do something - throw an exception similar to how its done in the python sdk?
created_at: 2025-03-05T10:22:16Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,202 @@
+# Copyright 2025 DataRobot, Inc. and its affiliates.
+#
+# All rights reserved.
+#
+# This is proprietary source code of DataRobot, Inc. and its affiliates.
+#
+# Released under the terms of DataRobot Tool and Utility Agreement.
+from typing import Any
+from typing import List
+from typing import Optional
+
+import datarobot as dr
+from airflow.exceptions import AirflowFailException
+from airflow.utils.context import Context
+from pandas import DataFrame
+
+from datarobot_provider.operators.base_datarobot_operator import BaseDatarobotOperator
+
+
+class GeneratePredictionsOperator(BaseDatarobotOperator):
+    """
+    Starts predictions generation job for a given model and dataset.
+
+    Args:
+        project_id (str): DataRobot project ID.
+        model_id (str): DataRobot model ID.
+        dataset_id (Optional[str]): DataRobot dataset ID. Defaults to None.
+        file_path (Optional[str]): File path to the dataset. Defaults to None.
+
+    Returns:
+        str: PredictJob ID.
+
+    Raises:
+        AirflowFailException: If neither dataset_id nor file_path is provided.
+        AirflowFailException: If both dataset_id and file_path are provided.
+    """
+
+    template_fields: List[str] = ["project_id", "model_id", "dataset_id", "file_path"]
+
+    def __init__(
+        self,
+        *,
+        project_id: str,
+        model_id: str,
+        dataset_id: Optional[str] = None,
+        file_path: Optional[str] = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.model_id = model_id
+        self.dataset_id = dataset_id
+        self.file_path = file_path
+
+    def validate(self) -> None:
+        if not self.dataset_id and not self.file_path:
+            raise AirflowFailException("Either dataset_id or file_path must be provided")
+        if self.dataset_id and self.file_path:
+            raise AirflowFailException("Both dataset_id and file_path cannot be provided")
+
+    def execute(self, context: Context) -> str:
+        model = dr.Model.get(
+            project=self.project_id,
+            model_id=self.model_id,
+        )
+
+        if self.dataset_id:
+            self.log.info(
+                f"Requesting predictions for model {self.model_id} using dataset {self.dataset_id}"
+            )
+            predict_job = model.request_predictions(dataset_id=self.dataset_id)
+        else:
+            self.log.info(
+                f"Requesting predictions for model {self.model_id} using file {self.file_path}"
+            )
+            predict_job = model.request_predictions(file_path=self.file_path)
+
+        self.log.info(f"Model predictions requested, predict_job_id={predict_job.id}")
+        return predict_job.id
+
+
+class GetPredictionsOperator(BaseDatarobotOperator):
+    """
+    Gets predictions from a PredictJob.
+
+    Args:
+        project_id (str): DataRobot project ID.
+        predict_job_id (str): PredictJob ID.
+
+    Returns:
+        DataFrame: Predictions DataFrame.
+
+    Raises:
+        AirflowFailException: If predict_job_id or project_id is not provided.
+    """
+
+    template_fields: List[str] = ["project_id", "predict_job_id"]
+
+    def __init__(self, *, project_id: str, predict_job_id: str, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.predict_job_id = predict_job_id
+
+    def validate(self) -> None:
+        if not self.predict_job_id or not self.project_id:
+            raise AirflowFailException("predict_job_id and project_id must be provided")
+
+    def execute(self, context: Context) -> DataFrame:
+        predictions = dr.PredictJob.get_predictions(
+            project_id=self.project_id,
+            predict_job_id=self.predict_job_id,
+        )
+
+        self.log.info(f"Predictions retrieved, predict_job_id={self.predict_job_id}")
+
+        ## NOTE: Can we even return a DataFrame here?
+        return predictions
+
+
+class GetPredictJobOperator(BaseDatarobotOperator):
+    """
+    Gets a PredictJob.
+
+    Args:
+        project_id (str): DataRobot project ID.
+        predict_job_id (str): PredictJob ID.
+
+    Returns:
+        dr.PredictJob: PredictJob object.
+
+    Raises:
+        AirflowFailException: If predict_job_id or project_id is not provided.
+        AirflowFailException: If PredictJob has already finished.
+    """
+
+    template_fields: List[str] = ["project_id", "predict_job_id"]
+
+    def __init__(self, *, project_id: str, predict_job_id: str, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.predict_job_id = predict_job_id
+
+    def validate(self) -> None:
+        if not self.predict_job_id or not self.project_id:
+            raise AirflowFailException("predict_job_id and project_id must be provided")
+
+    def execute(self, context: Context) -> dr.PredictJob:
+        try:
+            predict_job = dr.PredictJob.get(
+                project_id=self.project_id,
+                predict_job_id=self.predict_job_id,
+            )
+        except dr.errors.PendingJobFinished:
+            # NOTE: Is this the correct exception to raise in this case?
+            raise AirflowFailException(
+                "PredictJob has already finished, please use PredictJobGetPredictionsOperator to retrieve predictions"
+            ) from dr.errors.PendingJobFinished
+
+        self.log.info(f"PredictJob retrieved, predict_job_id={self.predict_job_id}")
+
+        # NOTE: Can we return the PredictJob object here?
+        return predict_job
path: datarobot_provider/operators/predictions.py
user: login: kessler-frost
body: > So, we are inputting the PredictJob id, and returning the same PredictJob id then?
> Seems a bit weird to return the same thing the user input, without doing much. Is the point that in case it finished running then we do something - throw an exception similar to how its done in the python sdk?

(Quoting my reply from another comment)
created_at: 2025-03-05T10:23:17Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,202 @@
+# Copyright 2025 DataRobot, Inc. and its affiliates.
+#
+# All rights reserved.
+#
+# This is proprietary source code of DataRobot, Inc. and its affiliates.
+#
+# Released under the terms of DataRobot Tool and Utility Agreement.
+from typing import Any
+from typing import List
+from typing import Optional
+
+import datarobot as dr
+from airflow.exceptions import AirflowFailException
+from airflow.utils.context import Context
+from pandas import DataFrame
+
+from datarobot_provider.operators.base_datarobot_operator import BaseDatarobotOperator
+
+
+class GeneratePredictionsOperator(BaseDatarobotOperator):
+    """
+    Starts predictions generation job for a given model and dataset.
+
+    Args:
+        project_id (str): DataRobot project ID.
+        model_id (str): DataRobot model ID.
+        dataset_id (Optional[str]): DataRobot dataset ID. Defaults to None.
+        file_path (Optional[str]): File path to the dataset. Defaults to None.
+
+    Returns:
+        str: PredictJob ID.
+
+    Raises:
+        AirflowFailException: If neither dataset_id nor file_path is provided.
+        AirflowFailException: If both dataset_id and file_path are provided.
+    """
+
+    template_fields: List[str] = ["project_id", "model_id", "dataset_id", "file_path"]
+
+    def __init__(
+        self,
+        *,
+        project_id: str,
+        model_id: str,
+        dataset_id: Optional[str] = None,
+        file_path: Optional[str] = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.model_id = model_id
+        self.dataset_id = dataset_id
+        self.file_path = file_path
+
+    def validate(self) -> None:
+        if not self.dataset_id and not self.file_path:
+            raise AirflowFailException("Either dataset_id or file_path must be provided")
+        if self.dataset_id and self.file_path:
+            raise AirflowFailException("Both dataset_id and file_path cannot be provided")
+
+    def execute(self, context: Context) -> str:
+        model = dr.Model.get(
+            project=self.project_id,
+            model_id=self.model_id,
+        )
+
+        if self.dataset_id:
+            self.log.info(
+                f"Requesting predictions for model {self.model_id} using dataset {self.dataset_id}"
+            )
+            predict_job = model.request_predictions(dataset_id=self.dataset_id)
+        else:
+            self.log.info(
+                f"Requesting predictions for model {self.model_id} using file {self.file_path}"
+            )
+            predict_job = model.request_predictions(file_path=self.file_path)
+
+        self.log.info(f"Model predictions requested, predict_job_id={predict_job.id}")
+        return predict_job.id
+
+
+class GetPredictionsOperator(BaseDatarobotOperator):
+    """
+    Gets predictions from a PredictJob.
+
+    Args:
+        project_id (str): DataRobot project ID.
+        predict_job_id (str): PredictJob ID.
+
+    Returns:
+        DataFrame: Predictions DataFrame.
+
+    Raises:
+        AirflowFailException: If predict_job_id or project_id is not provided.
+    """
+
+    template_fields: List[str] = ["project_id", "predict_job_id"]
+
+    def __init__(self, *, project_id: str, predict_job_id: str, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.predict_job_id = predict_job_id
+
+    def validate(self) -> None:
+        if not self.predict_job_id or not self.project_id:
+            raise AirflowFailException("predict_job_id and project_id must be provided")
+
+    def execute(self, context: Context) -> DataFrame:
+        predictions = dr.PredictJob.get_predictions(
+            project_id=self.project_id,
+            predict_job_id=self.predict_job_id,
+        )
+
+        self.log.info(f"Predictions retrieved, predict_job_id={self.predict_job_id}")
+
+        ## NOTE: Can we even return a DataFrame here?
+        return predictions
+
+
+class GetPredictJobOperator(BaseDatarobotOperator):
+    """
+    Gets a PredictJob.
+
+    Args:
+        project_id (str): DataRobot project ID.
+        predict_job_id (str): PredictJob ID.
+
+    Returns:
+        dr.PredictJob: PredictJob object.
+
+    Raises:
+        AirflowFailException: If predict_job_id or project_id is not provided.
+        AirflowFailException: If PredictJob has already finished.
+    """
+
+    template_fields: List[str] = ["project_id", "predict_job_id"]
+
+    def __init__(self, *, project_id: str, predict_job_id: str, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.predict_job_id = predict_job_id
+
+    def validate(self) -> None:
+        if not self.predict_job_id or not self.project_id:
+            raise AirflowFailException("predict_job_id and project_id must be provided")
+
+    def execute(self, context: Context) -> dr.PredictJob:
+        try:
+            predict_job = dr.PredictJob.get(
+                project_id=self.project_id,
+                predict_job_id=self.predict_job_id,
+            )
+        except dr.errors.PendingJobFinished:
+            # NOTE: Is this the correct exception to raise in this case?
+            raise AirflowFailException(
+                "PredictJob has already finished, please use PredictJobGetPredictionsOperator to retrieve predictions"
+            ) from dr.errors.PendingJobFinished
+
+        self.log.info(f"PredictJob retrieved, predict_job_id={self.predict_job_id}")
+
+        # NOTE: Can we return the PredictJob object here?
+        return predict_job
+
+
+class AwaitAndGetPredictionsOperator(BaseDatarobotOperator):
+    """
+    Waits and retrieves the predictions from a PredictJob.
+
+    Args:
+        project_id (str): DataRobot project ID.
+        predict_job_id (str): PredictJob ID.
+
+    Returns:
+        DataFrame: Predictions DataFrame.
+
+    Raises:
+        AirflowFailException: If predict_job_id or project_id is not provided.
+    """
+
+    template_fields: List[str] = ["project_id", "predict_job_id"]
+
+    def __init__(self, *, project_id: str, predict_job_id: str, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.predict_job_id = predict_job_id
+
+    def validate(self) -> None:
+        if not self.predict_job_id or not self.project_id:
+            raise AirflowFailException("predict_job_id and project_id must be provided")
+
+    def execute(self, context: Context) -> DataFrame:
+        predict_job = dr.PredictJob.get(
+            project_id=self.project_id,
+            predict_job_id=self.predict_job_id,
+        )
+
+        predictions = predict_job.get_result_when_complete()
+
+        self.log.info(f"Predictions waited for and retrieved, predict_job_id={self.predict_job_id}")
+
+        # NOTE: Again, can we return a DataFrame here?
+        return predictions
path: datarobot_provider/operators/predictions.py
user: login: kessler-frost
body: Sounds good, does it make sense to instead put the predictions in a dataset and return the id of that dataset? Or a similar s3 object?
created_at: 2025-03-05T10:24:23Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,202 @@
+# Copyright 2025 DataRobot, Inc. and its affiliates.
+#
+# All rights reserved.
+#
+# This is proprietary source code of DataRobot, Inc. and its affiliates.
+#
+# Released under the terms of DataRobot Tool and Utility Agreement.
+from typing import Any
+from typing import List
+from typing import Optional
+
+import datarobot as dr
+from airflow.exceptions import AirflowFailException
+from airflow.utils.context import Context
+from pandas import DataFrame
+
+from datarobot_provider.operators.base_datarobot_operator import BaseDatarobotOperator
+
+
+class GeneratePredictionsOperator(BaseDatarobotOperator):
path: datarobot_provider/operators/predictions.py
user: login: andrius-senulis
body: OK, thanks.
created_at: 2025-03-05T10:41:44Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,202 @@
+# Copyright 2025 DataRobot, Inc. and its affiliates.
+#
+# All rights reserved.
+#
+# This is proprietary source code of DataRobot, Inc. and its affiliates.
+#
+# Released under the terms of DataRobot Tool and Utility Agreement.
+from typing import Any
+from typing import List
+from typing import Optional
+
+import datarobot as dr
+from airflow.exceptions import AirflowFailException
+from airflow.utils.context import Context
+from pandas import DataFrame
+
+from datarobot_provider.operators.base_datarobot_operator import BaseDatarobotOperator
+
+
+class RequestPredictionsOperator(BaseDatarobotOperator):
+    """
+    Starts predictions generation job for a given model and dataset.
+
+    Args:
+        project_id (str): DataRobot project ID.
+        model_id (str): DataRobot model ID.
+        dataset_id (Optional[str]): DataRobot dataset ID. Defaults to None.
+        file_path (Optional[str]): File path to the dataset. Defaults to None.
+
+    Returns:
+        str: PredictJob ID.
+
+    Raises:
+        AirflowFailException: If neither dataset_id nor file_path is provided.
+        AirflowFailException: If both dataset_id and file_path are provided.
+    """
+
+    template_fields: List[str] = ["project_id", "model_id", "dataset_id", "file_path"]
+
+    def __init__(
+        self,
+        *,
+        project_id: str,
+        model_id: str,
+        dataset_id: Optional[str] = None,
+        file_path: Optional[str] = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.model_id = model_id
+        self.dataset_id = dataset_id
+        self.file_path = file_path
+
+    def validate(self) -> None:
+        if not self.dataset_id and not self.file_path:
+            raise AirflowFailException("Either dataset_id or file_path must be provided")
+        if self.dataset_id and self.file_path:
+            raise AirflowFailException("Both dataset_id and file_path cannot be provided")
+
+    def execute(self, context: Context) -> str:
+        model = dr.Model.get(
+            project=self.project_id,
+            model_id=self.model_id,
+        )
+
+        if self.dataset_id:
+            self.log.info(
+                f"Requesting predictions for model {self.model_id} using dataset {self.dataset_id}"
+            )
+            predict_job = model.request_predictions(dataset_id=self.dataset_id)
+        else:
+            self.log.info(
+                f"Requesting predictions for model {self.model_id} using file {self.file_path}"
+            )
+            predict_job = model.request_predictions(file_path=self.file_path)
+
+        self.log.info(f"Model predictions requested, predict_job_id={predict_job.id}")
+        return predict_job.id
+
+
+class GetPredictionsOperator(BaseDatarobotOperator):
+    """
+    Gets predictions from a PredictJob.
+
+    Args:
+        project_id (str): DataRobot project ID.
+        predict_job_id (str): PredictJob ID.
+
+    Returns:
+        DataFrame: Predictions DataFrame.
+
+    Raises:
+        AirflowFailException: If predict_job_id or project_id is not provided.
+    """
+
+    template_fields: List[str] = ["project_id", "predict_job_id"]
+
+    def __init__(self, *, project_id: str, predict_job_id: str, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.predict_job_id = predict_job_id
+
+    def validate(self) -> None:
+        if not self.predict_job_id or not self.project_id:
+            raise AirflowFailException("predict_job_id and project_id must be provided")
+
+    def execute(self, context: Context) -> DataFrame:
+        predictions = dr.PredictJob.get_predictions(
+            project_id=self.project_id,
+            predict_job_id=self.predict_job_id,
+        )
+
+        self.log.info(f"Predictions retrieved, predict_job_id={self.predict_job_id}")
+
+        ## NOTE: Can we even return a DataFrame here?
+        return predictions
+
+
+class GetPredictJobOperator(BaseDatarobotOperator):
+    """
+    Gets a PredictJob.
+
+    Args:
+        project_id (str): DataRobot project ID.
+        predict_job_id (str): PredictJob ID.
+
+    Returns:
+        dr.PredictJob: PredictJob object.
+
+    Raises:
+        AirflowFailException: If predict_job_id or project_id is not provided.
+        AirflowFailException: If PredictJob has already finished.
+    """
+
+    template_fields: List[str] = ["project_id", "predict_job_id"]
+
+    def __init__(self, *, project_id: str, predict_job_id: str, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.predict_job_id = predict_job_id
+
+    def validate(self) -> None:
+        if not self.predict_job_id or not self.project_id:
+            raise AirflowFailException("predict_job_id and project_id must be provided")
+
+    def execute(self, context: Context) -> dr.PredictJob:
+        try:
+            predict_job = dr.PredictJob.get(
+                project_id=self.project_id,
+                predict_job_id=self.predict_job_id,
+            )
+        except dr.errors.PendingJobFinished:
+            # NOTE: Is this the correct exception to raise in this case?
+            raise AirflowFailException(
+                "PredictJob has already finished, please use PredictJobGetPredictionsOperator to retrieve predictions"
+            ) from dr.errors.PendingJobFinished
+
+        self.log.info(f"PredictJob retrieved, predict_job_id={self.predict_job_id}")
+
+        # NOTE: Can we return the PredictJob object here?
+        return predict_job
+
+
+class AwaitAndGetPredictionsOperator(BaseDatarobotOperator):
path: datarobot_provider/operators/predictions.py
user: login: andrius-senulis
body: Unfortunately, I think this operator should not exist, since it's against the Airflow design. We should have a [Sensor](https://airflow.apache.org/docs/apache-airflow/stable/core-concepts/sensors.html) which just waits for completion instead. So the flow would be this:

`RequestPredictionsOperator -> WaitForPredictionsSensor -> GetPredictionsOperator`

The `GetPredictionsOperator` should store predictions somewhere, I'm not sure where. If we want to save in Catalog, then we should have a `SavePredictionsToCatalogOperator` or smth like that.
created_at: 2025-03-05T10:51:12Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,202 @@
+# Copyright 2025 DataRobot, Inc. and its affiliates.
+#
+# All rights reserved.
+#
+# This is proprietary source code of DataRobot, Inc. and its affiliates.
+#
+# Released under the terms of DataRobot Tool and Utility Agreement.
+from typing import Any
+from typing import List
+from typing import Optional
+
+import datarobot as dr
+from airflow.exceptions import AirflowFailException
+from airflow.utils.context import Context
+from pandas import DataFrame
+
+from datarobot_provider.operators.base_datarobot_operator import BaseDatarobotOperator
+
+
+class RequestPredictionsOperator(BaseDatarobotOperator):
+    """
+    Starts predictions generation job for a given model and dataset.
+
+    Args:
+        project_id (str): DataRobot project ID.
+        model_id (str): DataRobot model ID.
+        dataset_id (Optional[str]): DataRobot dataset ID. Defaults to None.
+        file_path (Optional[str]): File path to the dataset. Defaults to None.
+
+    Returns:
+        str: PredictJob ID.
+
+    Raises:
+        AirflowFailException: If neither dataset_id nor file_path is provided.
+        AirflowFailException: If both dataset_id and file_path are provided.
+    """
+
+    template_fields: List[str] = ["project_id", "model_id", "dataset_id", "file_path"]
+
+    def __init__(
+        self,
+        *,
+        project_id: str,
+        model_id: str,
+        dataset_id: Optional[str] = None,
+        file_path: Optional[str] = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.model_id = model_id
+        self.dataset_id = dataset_id
+        self.file_path = file_path
+
+    def validate(self) -> None:
+        if not self.dataset_id and not self.file_path:
+            raise AirflowFailException("Either dataset_id or file_path must be provided")
+        if self.dataset_id and self.file_path:
+            raise AirflowFailException("Both dataset_id and file_path cannot be provided")
+
+    def execute(self, context: Context) -> str:
+        model = dr.Model.get(
+            project=self.project_id,
+            model_id=self.model_id,
+        )
+
+        if self.dataset_id:
+            self.log.info(
+                f"Requesting predictions for model {self.model_id} using dataset {self.dataset_id}"
+            )
+            predict_job = model.request_predictions(dataset_id=self.dataset_id)
+        else:
+            self.log.info(
+                f"Requesting predictions for model {self.model_id} using file {self.file_path}"
+            )
+            predict_job = model.request_predictions(file_path=self.file_path)
+
+        self.log.info(f"Model predictions requested, predict_job_id={predict_job.id}")
+        return predict_job.id
+
+
+class GetPredictionsOperator(BaseDatarobotOperator):
+    """
+    Gets predictions from a PredictJob.
+
+    Args:
+        project_id (str): DataRobot project ID.
+        predict_job_id (str): PredictJob ID.
+
+    Returns:
+        DataFrame: Predictions DataFrame.
+
+    Raises:
+        AirflowFailException: If predict_job_id or project_id is not provided.
+    """
+
+    template_fields: List[str] = ["project_id", "predict_job_id"]
+
+    def __init__(self, *, project_id: str, predict_job_id: str, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.predict_job_id = predict_job_id
+
+    def validate(self) -> None:
+        if not self.predict_job_id or not self.project_id:
+            raise AirflowFailException("predict_job_id and project_id must be provided")
+
+    def execute(self, context: Context) -> DataFrame:
+        predictions = dr.PredictJob.get_predictions(
+            project_id=self.project_id,
+            predict_job_id=self.predict_job_id,
+        )
+
+        self.log.info(f"Predictions retrieved, predict_job_id={self.predict_job_id}")
+
+        ## NOTE: Can we even return a DataFrame here?
+        return predictions
+
+
+class GetPredictJobOperator(BaseDatarobotOperator):
path: datarobot_provider/operators/predictions.py
user: login: andrius-senulis
body: Unfortunately, I don't see the purpose of having this operator.
created_at: 2025-03-05T10:56:24Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,202 @@
+# Copyright 2025 DataRobot, Inc. and its affiliates.
+#
+# All rights reserved.
+#
+# This is proprietary source code of DataRobot, Inc. and its affiliates.
+#
+# Released under the terms of DataRobot Tool and Utility Agreement.
+from typing import Any
+from typing import List
+from typing import Optional
+
+import datarobot as dr
+from airflow.exceptions import AirflowFailException
+from airflow.utils.context import Context
+from pandas import DataFrame
+
+from datarobot_provider.operators.base_datarobot_operator import BaseDatarobotOperator
+
+
+class RequestPredictionsOperator(BaseDatarobotOperator):
+    """
+    Starts predictions generation job for a given model and dataset.
+
+    Args:
+        project_id (str): DataRobot project ID.
+        model_id (str): DataRobot model ID.
+        dataset_id (Optional[str]): DataRobot dataset ID. Defaults to None.
+        file_path (Optional[str]): File path to the dataset. Defaults to None.
+
+    Returns:
+        str: PredictJob ID.
+
+    Raises:
+        AirflowFailException: If neither dataset_id nor file_path is provided.
+        AirflowFailException: If both dataset_id and file_path are provided.
+    """
+
+    template_fields: List[str] = ["project_id", "model_id", "dataset_id", "file_path"]
+
+    def __init__(
+        self,
+        *,
+        project_id: str,
+        model_id: str,
+        dataset_id: Optional[str] = None,
+        file_path: Optional[str] = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.model_id = model_id
+        self.dataset_id = dataset_id
+        self.file_path = file_path
+
+    def validate(self) -> None:
+        if not self.dataset_id and not self.file_path:
+            raise AirflowFailException("Either dataset_id or file_path must be provided")
+        if self.dataset_id and self.file_path:
+            raise AirflowFailException("Both dataset_id and file_path cannot be provided")
+
+    def execute(self, context: Context) -> str:
+        model = dr.Model.get(
+            project=self.project_id,
+            model_id=self.model_id,
+        )
+
+        if self.dataset_id:
+            self.log.info(
+                f"Requesting predictions for model {self.model_id} using dataset {self.dataset_id}"
+            )
+            predict_job = model.request_predictions(dataset_id=self.dataset_id)
+        else:
+            self.log.info(
+                f"Requesting predictions for model {self.model_id} using file {self.file_path}"
+            )
+            predict_job = model.request_predictions(file_path=self.file_path)
+
+        self.log.info(f"Model predictions requested, predict_job_id={predict_job.id}")
+        return predict_job.id
+
+
+class GetPredictionsOperator(BaseDatarobotOperator):
+    """
+    Gets predictions from a PredictJob.
+
+    Args:
+        project_id (str): DataRobot project ID.
+        predict_job_id (str): PredictJob ID.
+
+    Returns:
+        DataFrame: Predictions DataFrame.
+
+    Raises:
+        AirflowFailException: If predict_job_id or project_id is not provided.
+    """
+
+    template_fields: List[str] = ["project_id", "predict_job_id"]
+
+    def __init__(self, *, project_id: str, predict_job_id: str, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.predict_job_id = predict_job_id
+
+    def validate(self) -> None:
+        if not self.predict_job_id or not self.project_id:
+            raise AirflowFailException("predict_job_id and project_id must be provided")
+
+    def execute(self, context: Context) -> DataFrame:
+        predictions = dr.PredictJob.get_predictions(
+            project_id=self.project_id,
+            predict_job_id=self.predict_job_id,
+        )
+
+        self.log.info(f"Predictions retrieved, predict_job_id={self.predict_job_id}")
+
+        ## NOTE: Can we even return a DataFrame here?
+        return predictions
+
+
+class GetPredictJobOperator(BaseDatarobotOperator):
+    """
+    Gets a PredictJob.
+
+    Args:
+        project_id (str): DataRobot project ID.
+        predict_job_id (str): PredictJob ID.
+
+    Returns:
+        dr.PredictJob: PredictJob object.
+
+    Raises:
+        AirflowFailException: If predict_job_id or project_id is not provided.
+        AirflowFailException: If PredictJob has already finished.
+    """
+
+    template_fields: List[str] = ["project_id", "predict_job_id"]
+
+    def __init__(self, *, project_id: str, predict_job_id: str, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.predict_job_id = predict_job_id
+
+    def validate(self) -> None:
+        if not self.predict_job_id or not self.project_id:
+            raise AirflowFailException("predict_job_id and project_id must be provided")
+
+    def execute(self, context: Context) -> dr.PredictJob:
+        try:
+            predict_job = dr.PredictJob.get(
+                project_id=self.project_id,
+                predict_job_id=self.predict_job_id,
+            )
+        except dr.errors.PendingJobFinished:
+            # NOTE: Is this the correct exception to raise in this case?
+            raise AirflowFailException(
+                "PredictJob has already finished, please use PredictJobGetPredictionsOperator to retrieve predictions"
+            ) from dr.errors.PendingJobFinished
+
+        self.log.info(f"PredictJob retrieved, predict_job_id={self.predict_job_id}")
+
+        # NOTE: Can we return the PredictJob object here?
+        return predict_job
+
+
+class AwaitAndGetPredictionsOperator(BaseDatarobotOperator):
path: datarobot_provider/operators/predictions.py
user: login: kessler-frost
body: Hmm, yeah I think that makes more sense. We can use the `PredictJob.wait_for_completion` in the `WaitForPredictionsSensor` sensor.

For the "predictions storage" question, I think we have 2 options depending on our use cases:

1. Have a standard place to store the predictions ALWAYS and return some sort of id of that - that makes it easier for users to not worry about managing another resource just for their predictions
2. Write separate operators to store the predictions in separate places like s3, catalog, or some other kind of store -> I don't think that's a good idea as it will get quite complicated very quickly - so unless there's a compelling use case, we should avoid doing this
created_at: 2025-03-05T11:01:36Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,202 @@
+# Copyright 2025 DataRobot, Inc. and its affiliates.
+#
+# All rights reserved.
+#
+# This is proprietary source code of DataRobot, Inc. and its affiliates.
+#
+# Released under the terms of DataRobot Tool and Utility Agreement.
+from typing import Any
+from typing import List
+from typing import Optional
+
+import datarobot as dr
+from airflow.exceptions import AirflowFailException
+from airflow.utils.context import Context
+from pandas import DataFrame
+
+from datarobot_provider.operators.base_datarobot_operator import BaseDatarobotOperator
+
+
+class RequestPredictionsOperator(BaseDatarobotOperator):
+    """
+    Starts predictions generation job for a given model and dataset.
+
+    Args:
+        project_id (str): DataRobot project ID.
+        model_id (str): DataRobot model ID.
+        dataset_id (Optional[str]): DataRobot dataset ID. Defaults to None.
+        file_path (Optional[str]): File path to the dataset. Defaults to None.
+
+    Returns:
+        str: PredictJob ID.
+
+    Raises:
+        AirflowFailException: If neither dataset_id nor file_path is provided.
+        AirflowFailException: If both dataset_id and file_path are provided.
+    """
+
+    template_fields: List[str] = ["project_id", "model_id", "dataset_id", "file_path"]
+
+    def __init__(
+        self,
+        *,
+        project_id: str,
+        model_id: str,
+        dataset_id: Optional[str] = None,
+        file_path: Optional[str] = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.model_id = model_id
+        self.dataset_id = dataset_id
+        self.file_path = file_path
+
+    def validate(self) -> None:
+        if not self.dataset_id and not self.file_path:
+            raise AirflowFailException("Either dataset_id or file_path must be provided")
+        if self.dataset_id and self.file_path:
+            raise AirflowFailException("Both dataset_id and file_path cannot be provided")
+
+    def execute(self, context: Context) -> str:
+        model = dr.Model.get(
+            project=self.project_id,
+            model_id=self.model_id,
+        )
+
+        if self.dataset_id:
+            self.log.info(
+                f"Requesting predictions for model {self.model_id} using dataset {self.dataset_id}"
+            )
+            predict_job = model.request_predictions(dataset_id=self.dataset_id)
+        else:
+            self.log.info(
+                f"Requesting predictions for model {self.model_id} using file {self.file_path}"
+            )
+            predict_job = model.request_predictions(file_path=self.file_path)
+
+        self.log.info(f"Model predictions requested, predict_job_id={predict_job.id}")
+        return predict_job.id
+
+
+class GetPredictionsOperator(BaseDatarobotOperator):
+    """
+    Gets predictions from a PredictJob.
+
+    Args:
+        project_id (str): DataRobot project ID.
+        predict_job_id (str): PredictJob ID.
+
+    Returns:
+        DataFrame: Predictions DataFrame.
+
+    Raises:
+        AirflowFailException: If predict_job_id or project_id is not provided.
+    """
+
+    template_fields: List[str] = ["project_id", "predict_job_id"]
+
+    def __init__(self, *, project_id: str, predict_job_id: str, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.predict_job_id = predict_job_id
+
+    def validate(self) -> None:
+        if not self.predict_job_id or not self.project_id:
+            raise AirflowFailException("predict_job_id and project_id must be provided")
+
+    def execute(self, context: Context) -> DataFrame:
+        predictions = dr.PredictJob.get_predictions(
+            project_id=self.project_id,
+            predict_job_id=self.predict_job_id,
+        )
+
+        self.log.info(f"Predictions retrieved, predict_job_id={self.predict_job_id}")
+
+        ## NOTE: Can we even return a DataFrame here?
+        return predictions
+
+
+class GetPredictJobOperator(BaseDatarobotOperator):
path: datarobot_provider/operators/predictions.py
user: login: kessler-frost
body: Yeah, if we can't use the returned object, then there's no reason to have this one.
created_at: 2025-03-05T11:03:08Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,202 @@
+# Copyright 2025 DataRobot, Inc. and its affiliates.
+#
+# All rights reserved.
+#
+# This is proprietary source code of DataRobot, Inc. and its affiliates.
+#
+# Released under the terms of DataRobot Tool and Utility Agreement.
+from typing import Any
+from typing import List
+from typing import Optional
+
+import datarobot as dr
+from airflow.exceptions import AirflowFailException
+from airflow.utils.context import Context
+from pandas import DataFrame
+
+from datarobot_provider.operators.base_datarobot_operator import BaseDatarobotOperator
+
+
+class RequestPredictionsOperator(BaseDatarobotOperator):
+    """
+    Starts predictions generation job for a given model and dataset.
+
+    Args:
+        project_id (str): DataRobot project ID.
+        model_id (str): DataRobot model ID.
+        dataset_id (Optional[str]): DataRobot dataset ID. Defaults to None.
+        file_path (Optional[str]): File path to the dataset. Defaults to None.
+
+    Returns:
+        str: PredictJob ID.
+
+    Raises:
+        AirflowFailException: If neither dataset_id nor file_path is provided.
+        AirflowFailException: If both dataset_id and file_path are provided.
+    """
+
+    template_fields: List[str] = ["project_id", "model_id", "dataset_id", "file_path"]
+
+    def __init__(
+        self,
+        *,
+        project_id: str,
+        model_id: str,
+        dataset_id: Optional[str] = None,
+        file_path: Optional[str] = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.model_id = model_id
+        self.dataset_id = dataset_id
+        self.file_path = file_path
+
+    def validate(self) -> None:
+        if not self.dataset_id and not self.file_path:
+            raise AirflowFailException("Either dataset_id or file_path must be provided")
+        if self.dataset_id and self.file_path:
+            raise AirflowFailException("Both dataset_id and file_path cannot be provided")
+
+    def execute(self, context: Context) -> str:
+        model = dr.Model.get(
+            project=self.project_id,
+            model_id=self.model_id,
+        )
+
+        if self.dataset_id:
+            self.log.info(
+                f"Requesting predictions for model {self.model_id} using dataset {self.dataset_id}"
+            )
+            predict_job = model.request_predictions(dataset_id=self.dataset_id)
+        else:
+            self.log.info(
+                f"Requesting predictions for model {self.model_id} using file {self.file_path}"
+            )
+            predict_job = model.request_predictions(file_path=self.file_path)
+
+        self.log.info(f"Model predictions requested, predict_job_id={predict_job.id}")
+        return predict_job.id
+
+
+class GetPredictionsOperator(BaseDatarobotOperator):
+    """
+    Gets predictions from a PredictJob.
+
+    Args:
+        project_id (str): DataRobot project ID.
+        predict_job_id (str): PredictJob ID.
+
+    Returns:
+        DataFrame: Predictions DataFrame.
+
+    Raises:
+        AirflowFailException: If predict_job_id or project_id is not provided.
+    """
+
+    template_fields: List[str] = ["project_id", "predict_job_id"]
+
+    def __init__(self, *, project_id: str, predict_job_id: str, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.predict_job_id = predict_job_id
+
+    def validate(self) -> None:
+        if not self.predict_job_id or not self.project_id:
+            raise AirflowFailException("predict_job_id and project_id must be provided")
+
+    def execute(self, context: Context) -> DataFrame:
+        predictions = dr.PredictJob.get_predictions(
+            project_id=self.project_id,
+            predict_job_id=self.predict_job_id,
+        )
+
+        self.log.info(f"Predictions retrieved, predict_job_id={self.predict_job_id}")
+
+        ## NOTE: Can we even return a DataFrame here?
+        return predictions
+
+
+class GetPredictJobOperator(BaseDatarobotOperator):
+    """
+    Gets a PredictJob.
+
+    Args:
+        project_id (str): DataRobot project ID.
+        predict_job_id (str): PredictJob ID.
+
+    Returns:
+        dr.PredictJob: PredictJob object.
+
+    Raises:
+        AirflowFailException: If predict_job_id or project_id is not provided.
+        AirflowFailException: If PredictJob has already finished.
+    """
+
+    template_fields: List[str] = ["project_id", "predict_job_id"]
+
+    def __init__(self, *, project_id: str, predict_job_id: str, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.predict_job_id = predict_job_id
+
+    def validate(self) -> None:
+        if not self.predict_job_id or not self.project_id:
+            raise AirflowFailException("predict_job_id and project_id must be provided")
+
+    def execute(self, context: Context) -> dr.PredictJob:
+        try:
+            predict_job = dr.PredictJob.get(
+                project_id=self.project_id,
+                predict_job_id=self.predict_job_id,
+            )
+        except dr.errors.PendingJobFinished:
+            # NOTE: Is this the correct exception to raise in this case?
+            raise AirflowFailException(
+                "PredictJob has already finished, please use PredictJobGetPredictionsOperator to retrieve predictions"
+            ) from dr.errors.PendingJobFinished
+
+        self.log.info(f"PredictJob retrieved, predict_job_id={self.predict_job_id}")
+
+        # NOTE: Can we return the PredictJob object here?
+        return predict_job
+
+
+class AwaitAndGetPredictionsOperator(BaseDatarobotOperator):
path: datarobot_provider/operators/predictions.py
user: login: kessler-frost
body: I'm gonna go ahead and create that sensor and disassemble this operator.
created_at: 2025-03-05T11:15:06Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,202 @@
+# Copyright 2025 DataRobot, Inc. and its affiliates.
+#
+# All rights reserved.
+#
+# This is proprietary source code of DataRobot, Inc. and its affiliates.
+#
+# Released under the terms of DataRobot Tool and Utility Agreement.
+from typing import Any
+from typing import List
+from typing import Optional
+
+import datarobot as dr
+from airflow.exceptions import AirflowFailException
+from airflow.utils.context import Context
+from pandas import DataFrame
+
+from datarobot_provider.operators.base_datarobot_operator import BaseDatarobotOperator
+
+
+class GeneratePredictionsOperator(BaseDatarobotOperator):
+    """
+    Starts predictions generation job for a given model and dataset.
+
+    Args:
+        project_id (str): DataRobot project ID.
+        model_id (str): DataRobot model ID.
+        dataset_id (Optional[str]): DataRobot dataset ID. Defaults to None.
+        file_path (Optional[str]): File path to the dataset. Defaults to None.
+
+    Returns:
+        str: PredictJob ID.
+
+    Raises:
+        AirflowFailException: If neither dataset_id nor file_path is provided.
+        AirflowFailException: If both dataset_id and file_path are provided.
+    """
+
+    template_fields: List[str] = ["project_id", "model_id", "dataset_id", "file_path"]
+
+    def __init__(
+        self,
+        *,
+        project_id: str,
+        model_id: str,
+        dataset_id: Optional[str] = None,
+        file_path: Optional[str] = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.model_id = model_id
+        self.dataset_id = dataset_id
+        self.file_path = file_path
+
+    def validate(self) -> None:
+        if not self.dataset_id and not self.file_path:
+            raise AirflowFailException("Either dataset_id or file_path must be provided")
+        if self.dataset_id and self.file_path:
+            raise AirflowFailException("Both dataset_id and file_path cannot be provided")
+
+    def execute(self, context: Context) -> str:
+        model = dr.Model.get(
+            project=self.project_id,
+            model_id=self.model_id,
+        )
+
+        if self.dataset_id:
+            self.log.info(
+                f"Requesting predictions for model {self.model_id} using dataset {self.dataset_id}"
+            )
+            predict_job = model.request_predictions(dataset_id=self.dataset_id)
+        else:
+            self.log.info(
+                f"Requesting predictions for model {self.model_id} using file {self.file_path}"
+            )
+            predict_job = model.request_predictions(file_path=self.file_path)
+
+        self.log.info(f"Model predictions requested, predict_job_id={predict_job.id}")
+        return predict_job.id
+
+
+class GetPredictionsOperator(BaseDatarobotOperator):
+    """
+    Gets predictions from a PredictJob.
+
+    Args:
+        project_id (str): DataRobot project ID.
+        predict_job_id (str): PredictJob ID.
+
+    Returns:
+        DataFrame: Predictions DataFrame.
+
+    Raises:
+        AirflowFailException: If predict_job_id or project_id is not provided.
+    """
+
+    template_fields: List[str] = ["project_id", "predict_job_id"]
+
+    def __init__(self, *, project_id: str, predict_job_id: str, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.predict_job_id = predict_job_id
+
+    def validate(self) -> None:
+        if not self.predict_job_id or not self.project_id:
+            raise AirflowFailException("predict_job_id and project_id must be provided")
+
+    def execute(self, context: Context) -> DataFrame:
+        predictions = dr.PredictJob.get_predictions(
+            project_id=self.project_id,
+            predict_job_id=self.predict_job_id,
+        )
+
+        self.log.info(f"Predictions retrieved, predict_job_id={self.predict_job_id}")
+
+        ## NOTE: Can we even return a DataFrame here?
+        return predictions
+
+
+class GetPredictJobOperator(BaseDatarobotOperator):
+    """
+    Gets a PredictJob.
+
+    Args:
+        project_id (str): DataRobot project ID.
+        predict_job_id (str): PredictJob ID.
+
+    Returns:
+        dr.PredictJob: PredictJob object.
+
+    Raises:
+        AirflowFailException: If predict_job_id or project_id is not provided.
+        AirflowFailException: If PredictJob has already finished.
+    """
+
+    template_fields: List[str] = ["project_id", "predict_job_id"]
+
+    def __init__(self, *, project_id: str, predict_job_id: str, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.predict_job_id = predict_job_id
+
+    def validate(self) -> None:
+        if not self.predict_job_id or not self.project_id:
+            raise AirflowFailException("predict_job_id and project_id must be provided")
+
+    def execute(self, context: Context) -> dr.PredictJob:
+        try:
+            predict_job = dr.PredictJob.get(
+                project_id=self.project_id,
+                predict_job_id=self.predict_job_id,
+            )
+        except dr.errors.PendingJobFinished:
path: datarobot_provider/operators/predictions.py
user: login: kessler-frost
body: removed this operator ✅ 
created_at: 2025-03-05T11:50:22Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,202 @@
+# Copyright 2025 DataRobot, Inc. and its affiliates.
+#
+# All rights reserved.
+#
+# This is proprietary source code of DataRobot, Inc. and its affiliates.
+#
+# Released under the terms of DataRobot Tool and Utility Agreement.
+from typing import Any
+from typing import List
+from typing import Optional
+
+import datarobot as dr
+from airflow.exceptions import AirflowFailException
+from airflow.utils.context import Context
+from pandas import DataFrame
+
+from datarobot_provider.operators.base_datarobot_operator import BaseDatarobotOperator
+
+
+class RequestPredictionsOperator(BaseDatarobotOperator):
+    """
+    Starts predictions generation job for a given model and dataset.
+
+    Args:
+        project_id (str): DataRobot project ID.
+        model_id (str): DataRobot model ID.
+        dataset_id (Optional[str]): DataRobot dataset ID. Defaults to None.
+        file_path (Optional[str]): File path to the dataset. Defaults to None.
+
+    Returns:
+        str: PredictJob ID.
+
+    Raises:
+        AirflowFailException: If neither dataset_id nor file_path is provided.
+        AirflowFailException: If both dataset_id and file_path are provided.
+    """
+
+    template_fields: List[str] = ["project_id", "model_id", "dataset_id", "file_path"]
+
+    def __init__(
+        self,
+        *,
+        project_id: str,
+        model_id: str,
+        dataset_id: Optional[str] = None,
+        file_path: Optional[str] = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.model_id = model_id
+        self.dataset_id = dataset_id
+        self.file_path = file_path
+
+    def validate(self) -> None:
+        if not self.dataset_id and not self.file_path:
+            raise AirflowFailException("Either dataset_id or file_path must be provided")
+        if self.dataset_id and self.file_path:
+            raise AirflowFailException("Both dataset_id and file_path cannot be provided")
+
+    def execute(self, context: Context) -> str:
+        model = dr.Model.get(
+            project=self.project_id,
+            model_id=self.model_id,
+        )
+
+        if self.dataset_id:
+            self.log.info(
+                f"Requesting predictions for model {self.model_id} using dataset {self.dataset_id}"
+            )
+            predict_job = model.request_predictions(dataset_id=self.dataset_id)
+        else:
+            self.log.info(
+                f"Requesting predictions for model {self.model_id} using file {self.file_path}"
+            )
+            predict_job = model.request_predictions(file_path=self.file_path)
+
+        self.log.info(f"Model predictions requested, predict_job_id={predict_job.id}")
+        return predict_job.id
+
+
+class GetPredictionsOperator(BaseDatarobotOperator):
+    """
+    Gets predictions from a PredictJob.
+
+    Args:
+        project_id (str): DataRobot project ID.
+        predict_job_id (str): PredictJob ID.
+
+    Returns:
+        DataFrame: Predictions DataFrame.
+
+    Raises:
+        AirflowFailException: If predict_job_id or project_id is not provided.
+    """
+
+    template_fields: List[str] = ["project_id", "predict_job_id"]
+
+    def __init__(self, *, project_id: str, predict_job_id: str, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.predict_job_id = predict_job_id
+
+    def validate(self) -> None:
+        if not self.predict_job_id or not self.project_id:
+            raise AirflowFailException("predict_job_id and project_id must be provided")
+
+    def execute(self, context: Context) -> DataFrame:
+        predictions = dr.PredictJob.get_predictions(
+            project_id=self.project_id,
+            predict_job_id=self.predict_job_id,
+        )
+
+        self.log.info(f"Predictions retrieved, predict_job_id={self.predict_job_id}")
+
+        ## NOTE: Can we even return a DataFrame here?
+        return predictions
+
+
+class GetPredictJobOperator(BaseDatarobotOperator):
+    """
+    Gets a PredictJob.
+
+    Args:
+        project_id (str): DataRobot project ID.
+        predict_job_id (str): PredictJob ID.
+
+    Returns:
+        dr.PredictJob: PredictJob object.
+
+    Raises:
+        AirflowFailException: If predict_job_id or project_id is not provided.
+        AirflowFailException: If PredictJob has already finished.
+    """
+
+    template_fields: List[str] = ["project_id", "predict_job_id"]
+
+    def __init__(self, *, project_id: str, predict_job_id: str, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.predict_job_id = predict_job_id
+
+    def validate(self) -> None:
+        if not self.predict_job_id or not self.project_id:
+            raise AirflowFailException("predict_job_id and project_id must be provided")
+
+    def execute(self, context: Context) -> dr.PredictJob:
+        try:
+            predict_job = dr.PredictJob.get(
+                project_id=self.project_id,
+                predict_job_id=self.predict_job_id,
+            )
+        except dr.errors.PendingJobFinished:
+            # NOTE: Is this the correct exception to raise in this case?
+            raise AirflowFailException(
+                "PredictJob has already finished, please use PredictJobGetPredictionsOperator to retrieve predictions"
+            ) from dr.errors.PendingJobFinished
+
+        self.log.info(f"PredictJob retrieved, predict_job_id={self.predict_job_id}")
+
+        # NOTE: Can we return the PredictJob object here?
+        return predict_job
+
+
+class AwaitAndGetPredictionsOperator(BaseDatarobotOperator):
path: datarobot_provider/operators/predictions.py
user: login: kessler-frost
body: Done, added the sensor and removed this operator ✅ .
created_at: 2025-03-05T11:51:17Z
author_association: CONTRIBUTOR

PR comments:
user: login: kessler-frost
created_at: 2025-03-06T14:51:44Z
body: After discussion, we decided to do a "save" on the predictions' results instead of doing a "get". Thus, implementing the `SavePredictionsToDatasetOperator` operator, and the `WaitForPredictionsSensor` sensor.
author_association: CONTRIBUTOR
user: login: devexp-slackbot[bot]
created_at: 2025-03-06T14:59:10Z
body: <p>The Needs Review labels were added based on the following file changes.</p><p><details><summary>Team @datarobot/machine-learning-development (<a href="https://datarobot.slack.com/archives/predictive-ai">#predictive-ai</a>) was assigned because of changes in files:</summary><p><pre>datarobot_provider/operators/predictions.py
datarobot_provider/sensors/predictions.py
docs/autodoc_operators/operators_predictions.md
docs/autodoc_sensors/sensors_predictions.md</pre></p></details></p><p>If you think that there are some issues with ownership, please discuss with C&A domain at #sdtk slack channel and create PR to update DRCODEOWNERS\CODEOWNERS file.</p>
author_association: NONE