[DM-17022] Default use case #144

demidov91 · 2025-02-07T10:50:33Z

This repository is public. Do not put here any private DataRobot or customer's data: code, datasets, model artifacts, .etc.

Summary

Introduce a notion of a default Use Case for a DAG.
Add a BaseUseCaseEntityOperator class
Rename CreateUseCaseOperator into GetOrCreateUseCaseOperator

Rationale

There are quite many DataRobot assets which have to belong to a UseCase. Define a default one in the very beginning of your DAG to avoid repetitive use_case_id parameters.

engprod-2 · 2025-02-07T14:36:53Z

The Needs Review labels were added based on the following file changes.

Team @datarobot/machine-learning-development (#machine-learning-dev) was assigned because of changes in files:

datarobot_provider/example_dags/hospital_readmissions_example.py
datarobot_provider/operators/base_datarobot_operator.py
datarobot_provider/operators/datarobot.py
tests/unit/conftest.py
tests/unit/operators/test_ai_catalog.py
tests/unit/operators/test_base_datarobot_operator.py
tests/unit/operators/test_datarobot.py

Team @datarobot/mldata (#mldata) was assigned because of changes in files:

datarobot_provider/operators/ai_catalog.py

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.

demidov91 · 2025-02-07T14:38:38Z

datarobot_provider/operators/base_datarobot_operator.py

+class BaseUseCaseEntityOperator(BaseDatarobotOperator):
+    """Base class for all operators which create a Use Case assets,
+    i.e. dataset, recipe, project, deployment, notebook, application, etc.
+
+    It introduces *use_case_id* template parameter and helper methods
+    to use a default Use Case if no Use Case was defined in the operator itself or a DAG context.
+    """
+
+    @classmethod
+    def __init_subclass__(cls, **kwargs):
+        super().__init_subclass__(**kwargs)
+        if "use_case_id" not in cls.template_fields:
+            cls.template_fields = (*cls.template_fields, "use_case_id")
+            cls.logger().info(
+                "use_case_id is implicitly added into %s template_fields list.", cls.__name__
+            )
+
+    def __init__(
+        self, *, use_case_id: Optional[str] = "{{ params.use_case_id | default('') }}", **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.use_case_id = use_case_id
+
+    def get_use_case(self, context: Context, required=False) -> Optional[dr.UseCase]:
+        """Get a `dr.UseCase` entity based on self.use_case_id or a default Use Case defined at runtime.
+        Raises an exception if no Use Case is defined and required=True
+        """
+        if use_case_id := self.get_use_case_id(context, required=required):
+            return dr.UseCase.get(use_case_id)
+
+        return None
+
+    def get_use_case_id(self, context: Context, required=False) -> Optional[str]:
+        """Get self.use_case_id or a default Use Case id defined at runtime.
+        Raises an exception if no Use Case id is defined and required=True"""
+        if use_case_id := (
+            self.use_case_id or self.xcom_pull(context, key=XCOM_DEFAULT_USE_CASE_ID)
+        ):
+            return use_case_id
+
+        if required:
+            raise AirflowException(
+                f"{self.__class__.__name__} requires a Use Case. Please define one of:\n"
+                "*use_case_id* parameter in the operator\n"
+                "*use_case_id* DAG context parameter\n"
+                "`GetOrCreateUseCaseOperator(..., set_default=True)` as one of the previous DAG tasks"
+            )
+
+        return None
+
+    def add_into_use_case(
+        self,
+        entity: Union[dr.Project, dr.Dataset, dr.models.Recipe, dr.Application],
+        *,
+        context: Context,
+    ):
+        """Add an *entity* into Use Case defined as self.use_case_id or a default DAG Use Case."""
+        if use_case := self.get_use_case(context):
+            use_case.add(entity)
+            self.log.info(
+                '%s is added into use case "%s"', entity.__class__.__name__, use_case.name
+            )
+
+        else:
+            self.log.info("The %s won't be added into any use case.", entity.__class__.__name__)


@c-h-russell-walker notebooks may find this base class useful.

Thank you!
Also I appreciate your review/comments on my draft PR - I've been sidelined by Fedramp 10.2 issues though so haven't gotten a chance to fully review/respond yet.

devexp-slackbot · 2025-02-07T16:06:50Z

Label Needs Review: MLData was removed because @bdrosen96 is part of Build Data & Features domain.

mjnitz02

I think this looks good. A few comments and questions as I'm trying to understand everything here.

mjnitz02 · 2025-02-07T17:27:54Z

datarobot_provider/operators/base_datarobot_operator.py

+
+        return None
+
+    def get_use_case_id(self, context: Context, required=False) -> Optional[str]:


Should required be a property of the Operator itself?

It seems like based on the operator type, required would be pre-determined, or at least defined through the init?

There are 3 ways to define a use case: an explicit parameter, a context variable and a default use case defined via xcom.

__init__ can handle only the first option.

Mandatory vs optional use_case_id could be really defined using requires_use_case: bool class field. At the moment PR uses an explicit argument in get_use_case_id() method. I'm ok with either option. I could switch to the class field either in this PR or in a followup if you find that option better.

mjnitz02 · 2025-02-07T17:30:13Z

datarobot_provider/operators/base_datarobot_operator.py

+        *,
+        context: Context,
+    ):
+        """Add an *entity* into Use Case defined as self.use_case_id or a default DAG Use Case."""


Is this mainly for adding datasets to a use case? I guess these are defined separately and then added to a use case? Sorry this might just be my misunderstanding. I guess like, a project, has to have a use case assigned to it during creation for it to be within the context of a use case?

It depends on the operator implantation. Yes, at the moment we use it for datasets. Nevertheless, it's applicable for projects or deployments, I guess

mjnitz02 · 2025-02-07T17:32:42Z

datarobot_provider/operators/base_datarobot_operator.py

+        self, *, use_case_id: Optional[str] = "{{ params.use_case_id | default('') }}", **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.use_case_id = use_case_id


I was thinking maybe this could auto-initialize, but after looking I guess context only exists in the scope of execute, so it wouldn't be present. It might be worth considering as some followup if we could sort of autowrap this so that the methods just use use_case_id = self.use_case_id and its either populated or None. I'm not 100% sure if thats actually possible to do in the context of how airflow works though. Just wanted to put the idea out there.

__init__ is not context-aware by design because operators are initialized much before the execution.

demidov91 added 5 commits February 7, 2025 10:54

Add a BaseUseCaseEntityOperator.

a822a7b

Update a DAG

04f7eb7

typecheck

0921e85

Define required use case

8633005

Brusing up

8985baa

demidov91 force-pushed the dzmitry/DM-17022-default-use-case branch from 8bff6a0 to 561f976 Compare February 7, 2025 14:36

Add unit tests

2f853c9

demidov91 force-pushed the dzmitry/DM-17022-default-use-case branch from 561f976 to 2f853c9 Compare February 7, 2025 14:36

demidov91 self-assigned this Feb 7, 2025

demidov91 added the 00 - Ready for Review label Feb 7, 2025

engprod-2 bot added Needs Review: MLData Needs Review: MLDev labels Feb 7, 2025

demidov91 commented Feb 7, 2025

View reviewed changes

demidov91 and others added 2 commits February 7, 2025 15:42

Add one more class to use a new base class.

2fa69d1

Merge branch 'main' into dzmitry/DM-17022-default-use-case

a6ef224

demidov91 requested a review from mjnitz02 February 7, 2025 15:55

bdrosen96 approved these changes Feb 7, 2025

View reviewed changes

devexp-slackbot bot removed the Needs Review: MLData label Feb 7, 2025

mjnitz02 removed the Needs Review: MLDev label Feb 7, 2025

mjnitz02 approved these changes Feb 7, 2025

View reviewed changes

demidov91 and others added 3 commits February 8, 2025 06:43

Merge branch 'main' into dzmitry/DM-17022-default-use-case

eceb245

Merge branch 'main' into dzmitry/DM-17022-default-use-case

251d13c

copyright

5d70621

demidov91 force-pushed the dzmitry/DM-17022-default-use-case branch from 8e7180b to 5d70621 Compare February 10, 2025 11:00

demidov91 merged commit b731e66 into main Feb 10, 2025
4 checks passed

demidov91 deleted the dzmitry/DM-17022-default-use-case branch February 10, 2025 11:04

svc-engprod-git1 removed the 00 - Ready for Review label Feb 10, 2025

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[DM-17022] Default use case #144

[DM-17022] Default use case #144

demidov91 commented Feb 7, 2025 •

edited

Loading

engprod-2 bot commented Feb 7, 2025

demidov91 Feb 7, 2025

c-h-russell-walker Feb 7, 2025

devexp-slackbot bot commented Feb 7, 2025

mjnitz02 left a comment

mjnitz02 Feb 7, 2025

mjnitz02 Feb 7, 2025

demidov91 Feb 7, 2025

mjnitz02 Feb 7, 2025

demidov91 Feb 7, 2025

mjnitz02 Feb 7, 2025

demidov91 Feb 7, 2025


		return None

		def get_use_case_id(self, context: Context, required=False) -> Optional[str]:

[DM-17022] Default use case #144

[DM-17022] Default use case #144

Conversation

demidov91 commented Feb 7, 2025 • edited Loading

This repository is public. Do not put here any private DataRobot or customer's data: code, datasets, model artifacts, .etc.

Summary

Rationale

engprod-2 bot commented Feb 7, 2025

Choose a reason for hiding this comment

Choose a reason for hiding this comment

devexp-slackbot bot commented Feb 7, 2025

mjnitz02 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

demidov91 commented Feb 7, 2025 •

edited

Loading