acryl-datahub 1.0.0.3rc4__py3-none-any.whl → 1.0.0.3rc5__py3-none-any.whl

{acryl_datahub-1.0.0.3rc4.dist-info → acryl_datahub-1.0.0.3rc5.dist-info}/RECORD

@@ -1,7 +1,7 @@
-acryl_datahub-1.0.0.3rc4.dist-info/licenses/LICENSE,sha256=9xNHpsD0uYF5ONzXsKDCuHHB-xbiCrSbueWXqrTNsxk,11365
+acryl_datahub-1.0.0.3rc5.dist-info/licenses/LICENSE,sha256=9xNHpsD0uYF5ONzXsKDCuHHB-xbiCrSbueWXqrTNsxk,11365
 datahub/__init__.py,sha256=aq_i5lVREmoLfYIqcx_pEQicO855YlhD19tWc1eZZNI,59
 datahub/__main__.py,sha256=pegIvQ9hzK7IhqVeUi1MeADSZ2QlP-D3K0OQdEg55RU,106
-datahub/_version.py,sha256=911KmAC1s4iZRnTeHRY8SXqgCBef_UX6nB2GAa1FqkE,323
+datahub/_version.py,sha256=QrGnvYk4Mo48c9uZMV5srbywzBX4Fcxp27GXbDul8F4,323
 datahub/entrypoints.py,sha256=2TYgHhs3sCxJlojIHjqfxzt3_ImPwPzq4vBtsUuMqu4,8885
 datahub/errors.py,sha256=BzKdcmYseHOt36zfjJXc17WNutFhp9Y23cU_L6cIkxc,612
 datahub/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -899,15 +899,17 @@ datahub/metadata/schemas/__init__.py,sha256=uvLNC3VyCkWA_v8e9FdA1leFf46NFKDD0Aaj
 datahub/pydantic/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datahub/pydantic/compat.py,sha256=TUEo4kSEeOWVAhV6LQtst1phrpVgGtK4uif4OI5vQ2M,1937
 datahub/sdk/__init__.py,sha256=QeutS6Th8K4E4ZxXuoGrmvahN6zA9Oh9asKk5mw9AIk,1670
-datahub/sdk/_all_entities.py,sha256=lCa2oeudmrKN7suOBIT5tidGJB8pgL3xD80DQlsYUKI,399
+datahub/sdk/_all_entities.py,sha256=inbLFv2T7dhZpGfBY5FPhCWbyE0P0G8umOt0Bc7V4XA,520
 datahub/sdk/_attribution.py,sha256=0Trh8steVd27GOr9MKCZeawbuDD2_q3GIsZlCtHqEUg,1321
-datahub/sdk/_shared.py,sha256=pHVKEJ50BoLw0fLLAm9zYsynNDN_bPI26qlj8nk2iyY,19582
+datahub/sdk/_shared.py,sha256=5L1IkihLc7Pd2x0ypDs96kZ8ecm6o0-UZEn0J1Sffqw,24808
 datahub/sdk/_utils.py,sha256=aGE665Su8SGtj2CRDiTaXNYrJ8ADBsS0m4DmaXw79b8,1027
 datahub/sdk/container.py,sha256=yw_vw9Jl1wOYNwMHxQHLz5ZvVQVDWWHi9CWBR3hOCd8,7547
 datahub/sdk/dataset.py,sha256=5LG4c_8bHeSPYrW88KNXRgiPD8frBjR0OBVrrwdquU4,29152
 datahub/sdk/entity.py,sha256=Q29AbpS58L4gD8ETwoNIwG-ouytz4c0MSSFi6-jLl_4,6742
-datahub/sdk/entity_client.py,sha256=Sxe6H6Vr_tqLJu5KW7MJfLWJ6mgh4mbsx7u7MOBpM64,5052
+datahub/sdk/entity_client.py,sha256=1AC9J7-jv3rD-MFEPz2PnFrT8nFkj_WO0M-4nyVOtQk,5319
 datahub/sdk/main_client.py,sha256=h2MKRhR-BO0zGCMhF7z2bTncX4hagKrAYwR3wTNTtzA,3666
+datahub/sdk/mlmodel.py,sha256=amS-hHg5tT7zAqEHG17kSA60Q7td2DFtO-W2rEfb2rY,10206
+datahub/sdk/mlmodelgroup.py,sha256=_7IkqkLVeyqYVEUHTVePSDLQyESsnwht5ca1lcMODAg,7842
 datahub/sdk/resolver_client.py,sha256=nKMAZJt2tRSGfKSzoREIh43PXqjM3umLiYkYHJjo1io,3243
 datahub/sdk/search_client.py,sha256=BJR5t7Ff2oDNOGLcSCp9YHzrGKbgOQr7T8XQKGEpucw,3437
 datahub/sdk/search_filters.py,sha256=BcMhvG5hGYAATtLPLz4WLRjKApX2oLYrrcGn-CG__ek,12901
@@ -1046,8 +1048,8 @@ datahub_provider/operators/datahub_assertion_operator.py,sha256=uvTQ-jk2F0sbqqxp
 datahub_provider/operators/datahub_assertion_sensor.py,sha256=lCBj_3x1cf5GMNpHdfkpHuyHfVxsm6ff5x2Z5iizcAo,140
 datahub_provider/operators/datahub_operation_operator.py,sha256=aevDp2FzX7FxGlXrR0khoHNbxbhKR2qPEX5e8O2Jyzw,174
 datahub_provider/operators/datahub_operation_sensor.py,sha256=8fcdVBCEPgqy1etTXgLoiHoJrRt_nzFZQMdSzHqSG7M,168
-acryl_datahub-1.0.0.3rc4.dist-info/METADATA,sha256=66uBY8gH_YmKeBmwbECGBpJhT7JiyrJXfFJv5wcFgQQ,176965
-acryl_datahub-1.0.0.3rc4.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-acryl_datahub-1.0.0.3rc4.dist-info/entry_points.txt,sha256=o3mDeJXSKhsy7XLkuogihraiabBdLn9HaizYXPrxmk0,9710
-acryl_datahub-1.0.0.3rc4.dist-info/top_level.txt,sha256=iLjSrLK5ox1YVYcglRUkcvfZPvKlobBWx7CTUXx8_GI,25
-acryl_datahub-1.0.0.3rc4.dist-info/RECORD,,
+acryl_datahub-1.0.0.3rc5.dist-info/METADATA,sha256=4yRDLa5bvQ1tAjJ2uzc4mKVmegeKOASmLpseMx9A1K8,176989
+acryl_datahub-1.0.0.3rc5.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+acryl_datahub-1.0.0.3rc5.dist-info/entry_points.txt,sha256=o3mDeJXSKhsy7XLkuogihraiabBdLn9HaizYXPrxmk0,9710
+acryl_datahub-1.0.0.3rc5.dist-info/top_level.txt,sha256=iLjSrLK5ox1YVYcglRUkcvfZPvKlobBWx7CTUXx8_GI,25
+acryl_datahub-1.0.0.3rc5.dist-info/RECORD,,

datahub/_version.py

@@ -1,6 +1,6 @@
 # Published at https://pypi.org/project/acryl-datahub/.
 __package_name__ = "acryl-datahub"
-__version__ = "1.0.0.3rc4"
+__version__ = "1.0.0.3rc5"
 
 
 def is_dev_mode() -> bool:

datahub/sdk/_all_entities.py

@@ -3,11 +3,15 @@ from typing import Dict, List, Type
 from datahub.sdk.container import Container
 from datahub.sdk.dataset import Dataset
 from datahub.sdk.entity import Entity
+from datahub.sdk.mlmodel import MLModel
+from datahub.sdk.mlmodelgroup import MLModelGroup
 
 # TODO: Is there a better way to declare this?
 ENTITY_CLASSES_LIST: List[Type[Entity]] = [
     Container,
     Dataset,
+    MLModel,
+    MLModelGroup,
 ]
 
 ENTITY_CLASSES: Dict[str, Type[Entity]] = {

datahub/sdk/_shared.py

@@ -5,6 +5,7 @@ from datetime import datetime
 from typing import (
     TYPE_CHECKING,
     Callable,
+    Dict,
     List,
     Optional,
     Sequence,
@@ -14,6 +15,7 @@ from typing import (
 
 from typing_extensions import TypeAlias, assert_never
 
+import datahub.emitter.mce_builder as builder
 import datahub.metadata.schema_classes as models
 from datahub.emitter.mce_builder import (
     make_ts_millis,
@@ -30,12 +32,14 @@ from datahub.metadata.urns import (
     DataJobUrn,
     DataPlatformInstanceUrn,
     DataPlatformUrn,
+    DataProcessInstanceUrn,
     DatasetUrn,
     DomainUrn,
     GlossaryTermUrn,
     OwnershipTypeUrn,
     TagUrn,
     Urn,
+    VersionSetUrn,
 )
 from datahub.sdk._utils import add_list_unique, remove_list_unique
 from datahub.sdk.entity import Entity
@@ -52,6 +56,36 @@ ActorUrn: TypeAlias = Union[CorpUserUrn, CorpGroupUrn]
 
 _DEFAULT_ACTOR_URN = CorpUserUrn("__ingestion").urn()
 
+TrainingMetricsInputType: TypeAlias = Union[
+    List[models.MLMetricClass], Dict[str, Optional[str]]
+]
+HyperParamsInputType: TypeAlias = Union[
+    List[models.MLHyperParamClass], Dict[str, Optional[str]]
+]
+MLTrainingJobInputType: TypeAlias = Union[Sequence[Union[str, DataProcessInstanceUrn]]]
+
+
+def convert_training_metrics(
+    metrics: TrainingMetricsInputType,
+) -> List[models.MLMetricClass]:
+    if isinstance(metrics, dict):
+        return [
+            models.MLMetricClass(name=name, value=str(value))
+            for name, value in metrics.items()
+        ]
+    return metrics
+
+
+def convert_hyper_params(
+    params: HyperParamsInputType,
+) -> List[models.MLHyperParamClass]:
+    if isinstance(params, dict):
+        return [
+            models.MLHyperParamClass(name=name, value=str(value))
+            for name, value in params.items()
+        ]
+    return params
+
 
 def make_time_stamp(ts: Optional[datetime]) -> Optional[models.TimeStampClass]:
     if ts is None:
@@ -578,3 +612,109 @@ class HasInstitutionalMemory(Entity):
             self._link_key,
             self._parse_link_association_class(link),
         )
+
+
+class HasVersion(Entity):
+    """Mixin for entities that have version properties."""
+
+    def _get_version_props(self) -> Optional[models.VersionPropertiesClass]:
+        return self._get_aspect(models.VersionPropertiesClass)
+
+    def _ensure_version_props(self) -> models.VersionPropertiesClass:
+        version_props = self._get_version_props()
+        if version_props is None:
+            guid_dict = {"urn": str(self.urn)}
+            version_set_urn = VersionSetUrn(
+                id=builder.datahub_guid(guid_dict), entity_type=self.urn.ENTITY_TYPE
+            )
+
+            version_props = models.VersionPropertiesClass(
+                versionSet=str(version_set_urn),
+                version=models.VersionTagClass(versionTag="0.1.0"),
+                sortId="0000000.1.0",
+            )
+            self._set_aspect(version_props)
+        return version_props
+
+    @property
+    def version(self) -> Optional[str]:
+        version_props = self._get_version_props()
+        if version_props and version_props.version:
+            return version_props.version.versionTag
+        return None
+
+    def set_version(self, version: str) -> None:
+        """Set the version of the entity."""
+        guid_dict = {"urn": str(self.urn)}
+        version_set_urn = VersionSetUrn(
+            id=builder.datahub_guid(guid_dict), entity_type=self.urn.ENTITY_TYPE
+        )
+
+        version_props = self._get_version_props()
+        if version_props is None:
+            # If no version properties exist, create a new one
+            version_props = models.VersionPropertiesClass(
+                version=models.VersionTagClass(versionTag=version),
+                versionSet=str(version_set_urn),
+                sortId=version.zfill(10),  # Pad with zeros for sorting
+            )
+        else:
+            # Update existing version properties
+            version_props.version = models.VersionTagClass(versionTag=version)
+            version_props.versionSet = str(version_set_urn)
+            version_props.sortId = version.zfill(10)
+
+        self._set_aspect(version_props)
+
+    @property
+    def version_aliases(self) -> List[str]:
+        version_props = self._get_version_props()
+        if version_props and version_props.aliases:
+            return [
+                alias.versionTag
+                for alias in version_props.aliases
+                if alias.versionTag is not None
+            ]
+        return []  # Return empty list instead of None
+
+    def set_version_aliases(self, aliases: List[str]) -> None:
+        version_props = self._get_aspect(models.VersionPropertiesClass)
+        if version_props:
+            version_props.aliases = [
+                models.VersionTagClass(versionTag=alias) for alias in aliases
+            ]
+        else:
+            # If no version properties exist, we need to create one with a default version
+            guid_dict = {"urn": str(self.urn)}
+            version_set_urn = VersionSetUrn(
+                id=builder.datahub_guid(guid_dict), entity_type=self.urn.ENTITY_TYPE
+            )
+            self._set_aspect(
+                models.VersionPropertiesClass(
+                    version=models.VersionTagClass(
+                        versionTag="0.1.0"
+                    ),  # Default version
+                    versionSet=str(version_set_urn),
+                    sortId="0000000.1.0",
+                    aliases=[
+                        models.VersionTagClass(versionTag=alias) for alias in aliases
+                    ],
+                )
+            )
+
+    def add_version_alias(self, alias: str) -> None:
+        if not alias:
+            raise ValueError("Alias cannot be empty")
+        version_props = self._ensure_version_props()
+        if version_props.aliases is None:
+            version_props.aliases = []
+        version_props.aliases.append(models.VersionTagClass(versionTag=alias))
+        self._set_aspect(version_props)
+
+    def remove_version_alias(self, alias: str) -> None:
+        version_props = self._get_version_props()
+        if version_props and version_props.aliases:
+            version_props.aliases = [
+                a for a in version_props.aliases if a.versionTag != alias
+            ]
+            self._set_aspect(version_props)

datahub/sdk/entity_client.py

@@ -11,6 +11,8 @@ from datahub.ingestion.graph.client import DataHubGraph
 from datahub.metadata.urns import (
     ContainerUrn,
     DatasetUrn,
+    MlModelGroupUrn,
+    MlModelUrn,
     Urn,
 )
 from datahub.sdk._all_entities import ENTITY_CLASSES
@@ -18,6 +20,8 @@ from datahub.sdk._shared import UrnOrStr
 from datahub.sdk.container import Container
 from datahub.sdk.dataset import Dataset
 from datahub.sdk.entity import Entity
+from datahub.sdk.mlmodel import MLModel
+from datahub.sdk.mlmodelgroup import MLModelGroup
 
 if TYPE_CHECKING:
     from datahub.sdk.main_client import DataHubClient
@@ -49,6 +53,10 @@ class EntityClient:
     @overload
     def get(self, urn: DatasetUrn) -> Dataset: ...
     @overload
+    def get(self, urn: MlModelUrn) -> MLModel: ...
+    @overload
+    def get(self, urn: MlModelGroupUrn) -> MLModelGroup: ...
+    @overload
     def get(self, urn: Union[Urn, str]) -> Entity: ...
     def get(self, urn: UrnOrStr) -> Entity:
         """Retrieve an entity by its urn.

datahub/sdk/mlmodel.py

@@ -0,0 +1,301 @@
+from __future__ import annotations
+
+from datetime import datetime
+from typing import Dict, List, Optional, Sequence, Type, Union
+
+from typing_extensions import Self
+
+from datahub.emitter.mce_builder import DEFAULT_ENV
+from datahub.metadata.schema_classes import (
+    AspectBag,
+    MLHyperParamClass,
+    MLMetricClass,
+    MLModelPropertiesClass,
+)
+from datahub.metadata.urns import (
+    DataProcessInstanceUrn,
+    MlModelGroupUrn,
+    MlModelUrn,
+    Urn,
+)
+from datahub.sdk._shared import (
+    DomainInputType,
+    HasDomain,
+    HasInstitutionalMemory,
+    HasOwnership,
+    HasPlatformInstance,
+    HasTags,
+    HasTerms,
+    HasVersion,
+    HyperParamsInputType,
+    LinksInputType,
+    MLTrainingJobInputType,
+    OwnersInputType,
+    TagsInputType,
+    TermsInputType,
+    TrainingMetricsInputType,
+    convert_hyper_params,
+    convert_training_metrics,
+    make_time_stamp,
+    parse_time_stamp,
+)
+from datahub.sdk.entity import Entity, ExtraAspectsType
+
+
+class MLModel(
+    HasPlatformInstance,
+    HasOwnership,
+    HasInstitutionalMemory,
+    HasTags,
+    HasTerms,
+    HasDomain,
+    HasVersion,
+    Entity,
+):
+    __slots__ = ()
+
+    @classmethod
+    def get_urn_type(cls) -> Type[MlModelUrn]:
+        return MlModelUrn
+
+    def __init__(
+        self,
+        id: str,
+        platform: str,
+        version: Optional[str] = None,
+        aliases: Optional[List[str]] = None,
+        platform_instance: Optional[str] = None,
+        env: str = DEFAULT_ENV,
+        name: Optional[str] = None,
+        description: Optional[str] = None,
+        training_metrics: Optional[TrainingMetricsInputType] = None,
+        hyper_params: Optional[HyperParamsInputType] = None,
+        external_url: Optional[str] = None,
+        custom_properties: Optional[Dict[str, str]] = None,
+        created: Optional[datetime] = None,
+        last_modified: Optional[datetime] = None,
+        owners: Optional[OwnersInputType] = None,
+        links: Optional[LinksInputType] = None,
+        tags: Optional[TagsInputType] = None,
+        terms: Optional[TermsInputType] = None,
+        domain: Optional[DomainInputType] = None,
+        model_group: Optional[Union[str, MlModelGroupUrn]] = None,
+        training_jobs: Optional[MLTrainingJobInputType] = None,
+        downstream_jobs: Optional[MLTrainingJobInputType] = None,
+        extra_aspects: ExtraAspectsType = None,
+    ):
+        urn = MlModelUrn(platform=platform, name=id, env=env)
+        super().__init__(urn)
+        self._set_extra_aspects(extra_aspects)
+
+        self._set_platform_instance(urn.platform, platform_instance)
+
+        self._ensure_model_props()
+
+        if version is not None:
+            self.set_version(version)
+        if name is not None:
+            self.set_name(name)
+        if aliases is not None:
+            self.set_version_aliases(aliases)
+        if description is not None:
+            self.set_description(description)
+        if training_metrics is not None:
+            self.set_training_metrics(training_metrics)
+        if hyper_params is not None:
+            self.set_hyper_params(hyper_params)
+        if external_url is not None:
+            self.set_external_url(external_url)
+        if custom_properties is not None:
+            self.set_custom_properties(custom_properties)
+        if created is not None:
+            self.set_created(created)
+        if last_modified is not None:
+            self.set_last_modified(last_modified)
+
+        if owners is not None:
+            self.set_owners(owners)
+        if links is not None:
+            self.set_links(links)
+        if tags is not None:
+            self.set_tags(tags)
+        if terms is not None:
+            self.set_terms(terms)
+        if domain is not None:
+            self.set_domain(domain)
+        if model_group is not None:
+            self.set_model_group(model_group)
+        if training_jobs is not None:
+            self.set_training_jobs(training_jobs)
+        if downstream_jobs is not None:
+            self.set_downstream_jobs(downstream_jobs)
+
+    @classmethod
+    def _new_from_graph(cls, urn: Urn, current_aspects: AspectBag) -> Self:
+        assert isinstance(urn, MlModelUrn)
+        entity = cls(
+            id=urn.name,
+            platform=urn.platform,
+            env=urn.env,
+        )
+        return entity._init_from_graph(current_aspects)
+
+    @property
+    def urn(self) -> MlModelUrn:
+        return self._urn  # type: ignore
+
+    def _ensure_model_props(
+        self,
+    ) -> MLModelPropertiesClass:
+        return self._setdefault_aspect(MLModelPropertiesClass())
+
+    @property
+    def name(self) -> Optional[str]:
+        return self._ensure_model_props().name
+
+    def set_name(self, name: str) -> None:
+        self._ensure_model_props().name = name
+
+    @property
+    def description(self) -> Optional[str]:
+        return self._ensure_model_props().description
+
+    def set_description(self, description: str) -> None:
+        self._ensure_model_props().description = description
+
+    @property
+    def external_url(self) -> Optional[str]:
+        return self._ensure_model_props().externalUrl
+
+    def set_external_url(self, external_url: str) -> None:
+        self._ensure_model_props().externalUrl = external_url
+
+    @property
+    def custom_properties(self) -> Optional[Dict[str, str]]:
+        return self._ensure_model_props().customProperties
+
+    def set_custom_properties(self, custom_properties: Dict[str, str]) -> None:
+        self._ensure_model_props().customProperties = custom_properties
+
+    @property
+    def created(self) -> Optional[datetime]:
+        return parse_time_stamp(self._ensure_model_props().created)
+
+    def set_created(self, created: datetime) -> None:
+        self._ensure_model_props().created = make_time_stamp(created)
+
+    @property
+    def last_modified(self) -> Optional[datetime]:
+        return parse_time_stamp(self._ensure_model_props().lastModified)
+
+    def set_last_modified(self, last_modified: datetime) -> None:
+        self._ensure_model_props().lastModified = make_time_stamp(last_modified)
+
+    @property
+    def training_metrics(self) -> Optional[List[MLMetricClass]]:
+        return self._ensure_model_props().trainingMetrics
+
+    def set_training_metrics(self, metrics: TrainingMetricsInputType) -> None:
+        self._ensure_model_props().trainingMetrics = convert_training_metrics(metrics)
+
+    def add_training_metrics(self, metrics: TrainingMetricsInputType) -> None:
+        props = self._ensure_model_props()
+        if props.trainingMetrics is None:
+            props.trainingMetrics = []
+        if isinstance(metrics, list):
+            props.trainingMetrics.extend(
+                [
+                    MLMetricClass(name=metric.name, value=metric.value)
+                    for metric in metrics
+                ]
+            )
+        else:
+            # For dictionary case, use the key as name and value as value
+            for name, value in metrics.items():
+                props.trainingMetrics.append(MLMetricClass(name=name, value=value))
+
+    @property
+    def hyper_params(self) -> Optional[List[MLHyperParamClass]]:
+        return self._ensure_model_props().hyperParams
+
+    def set_hyper_params(self, params: HyperParamsInputType) -> None:
+        self._ensure_model_props().hyperParams = convert_hyper_params(params)
+
+    def add_hyper_params(self, params: HyperParamsInputType) -> None:
+        props = self._ensure_model_props()
+        if props.hyperParams is None:
+            props.hyperParams = []
+        if isinstance(params, list):
+            props.hyperParams.extend(
+                [
+                    MLHyperParamClass(name=param.name, value=param.value)
+                    for param in params
+                ]
+            )
+        else:
+            # For dictionary case, iterate through key-value pairs
+            for name, value in params.items():
+                props.hyperParams.append(MLHyperParamClass(name=name, value=value))
+
+    @property
+    def model_group(self) -> Optional[str]:
+        props = self._ensure_model_props()
+        groups = props.groups
+        if groups is None or len(groups) == 0:
+            return None
+        return groups[0]
+
+    def set_model_group(self, group: Union[str, MlModelGroupUrn]) -> None:
+        self._ensure_model_props().groups = [str(group)]
+
+    @property
+    def training_jobs(self) -> Optional[List[str]]:
+        return self._ensure_model_props().trainingJobs
+
+    def set_training_jobs(self, training_jobs: MLTrainingJobInputType) -> None:
+        self._ensure_model_props().trainingJobs = [str(job) for job in training_jobs]
+
+    def add_training_job(
+        self, training_job: Union[str, DataProcessInstanceUrn]
+    ) -> None:
+        props = self._ensure_model_props()
+        if props.trainingJobs is None:
+            props.trainingJobs = []
+        props.trainingJobs.append(str(training_job))
+
+    def remove_training_job(
+        self, training_job: Union[str, DataProcessInstanceUrn]
+    ) -> None:
+        props = self._ensure_model_props()
+        if props.trainingJobs is not None:
+            job_str = str(training_job)
+            props.trainingJobs = [job for job in props.trainingJobs if job != job_str]
+
+    @property
+    def downstream_jobs(self) -> Optional[List[str]]:
+        return self._ensure_model_props().downstreamJobs
+
+    def set_downstream_jobs(
+        self, downstream_jobs: Sequence[Union[str, DataProcessInstanceUrn]]
+    ) -> None:
+        self._ensure_model_props().downstreamJobs = [
+            str(job) for job in downstream_jobs
+        ]
+
+    def add_downstream_job(
+        self, downstream_job: Union[str, DataProcessInstanceUrn]
+    ) -> None:
+        props = self._ensure_model_props()
+        if props.downstreamJobs is None:
+            props.downstreamJobs = []
+        props.downstreamJobs.append(str(downstream_job))
+
+    def remove_downstream_job(
+        self, downstream_job: Union[str, DataProcessInstanceUrn]
+    ) -> None:
+        props = self._ensure_model_props()
+        if props.downstreamJobs is not None:
+            job_str = str(downstream_job)
+            props.downstreamJobs = [
+                job for job in props.downstreamJobs if job != job_str
+            ]