frogml-core 0.0.1__py3-none-any.whl → 0.0.2__py3-none-any.whl

frogml_core/feature_store/_common/feature_set_utils.py

@@ -0,0 +1,261 @@
+import time
+from dataclasses import dataclass
+from typing import Dict, Optional
+
+from frogml_proto.qwak.build.v1.build_pb2 import BatchFeature as ProtoBatchFeature
+from frogml_proto.qwak.build.v1.build_pb2 import BatchFeatureV1 as ProtoBatchFeatureV1
+from frogml_proto.qwak.build.v1.build_pb2 import Feature as ProtoFeature
+from frogml_proto.qwak.build.v1.build_pb2 import (
+    StreamingAggregationFeature as ProtoStreamingAggregationFeature,
+)
+from frogml_proto.qwak.build.v1.build_pb2 import (
+    StreamingFeature as ProtoStreamingFeature,
+)
+from frogml_proto.qwak.build.v1.build_pb2 import (
+    StreamingFeatureV1 as ProtoStreamingFeatureV1,
+)
+from frogml_proto.qwak.feature_store.entities.entity_pb2 import (
+    EntitySpec as ProtoEntitySpec,
+)
+from frogml_proto.qwak.feature_store.features.feature_set_types_pb2 import (
+    BatchFeatureSet as ProtoBatchFeatureSet,
+)
+from frogml_proto.qwak.feature_store.features.feature_set_types_pb2 import (
+    BatchFeatureSetV1 as ProtoBatchFeatureSetV1,
+)
+from frogml_proto.qwak.feature_store.features.feature_set_types_pb2 import (
+    FeatureSetType as ProtoFeatureSetType,
+)
+from frogml_proto.qwak.feature_store.features.feature_set_types_pb2 import (
+    StreamingAggregationFeatureSet as ProtoStreamingAggregationFeatureSet,
+)
+from frogml_proto.qwak.feature_store.features.feature_set_types_pb2 import (
+    StreamingFeatureSetV1 as ProtoStreamingFeatureSetV1,
+)
+from frogml_core.clients.feature_store import FeatureRegistryClient
+from frogml_core.exceptions import FrogmlException
+from frogml_core.model.schema_entities import BaseFeature, Entity, FeatureStoreInput
+
+_INTERNAL_KEY_PREFIX: str = "_frogml_key_"
+
+
+@dataclass(unsafe_hash=True)
+class BatchFeature(BaseFeature):
+    entity: Optional[Entity] = None
+
+    def to_proto(self):
+        return ProtoFeature(
+            batch_feature=ProtoBatchFeature(
+                name=self.name, entity=self.entity.to_proto()
+            )
+        )
+
+
+@dataclass(unsafe_hash=True)
+class BatchFeatureV1(BaseFeature):
+    entity: Optional[Entity] = None
+
+    def to_proto(self):
+        return ProtoFeature(
+            batch_feature_v1=ProtoBatchFeatureV1(
+                name=self.name, entity=self.entity.to_proto()
+            )
+        )
+
+
+@dataclass(unsafe_hash=True)
+class StreamingFeature(BaseFeature):
+    entity: Entity
+
+    def to_proto(self):
+        return ProtoFeature(
+            streaming_feature=ProtoStreamingFeature(
+                name=self.name, entity=self.entity.to_proto()
+            )
+        )
+
+
+@dataclass(unsafe_hash=True)
+class StreamingFeatureV1(BaseFeature):
+    entity: Optional[Entity] = None
+
+    def to_proto(self):
+        return ProtoFeature(
+            streaming_feature_v1=ProtoStreamingFeatureV1(
+                name=self.name, entity=self.entity.to_proto()
+            )
+        )
+
+
+@dataclass(unsafe_hash=True)
+class StreamingAggregationFeature(BaseFeature):
+    entity: Optional[Entity] = None
+
+    def to_proto(self):
+        return ProtoFeature(
+            streaming_aggregation_feature=ProtoStreamingAggregationFeature(
+                name=self.name, entity=self.entity.to_proto()
+            )
+        )
+
+
+@dataclass()
+class FeatureSetInfo:
+    entity_spec: ProtoEntitySpec
+    feature_set_type: ProtoFeatureSetType
+    feature_version: int
+
+
+def get_feature_set_info(
+    feature_manager_client: FeatureRegistryClient, feature_set_name: str
+) -> FeatureSetInfo:
+    """
+    Get the entities by the feature set name and feature type
+    Args:
+        feature_manager_client: feature manager client for the grpc request
+        feature_set_name: the required feature set name
+    Returns: tuple of entity spec, type and feature version
+
+    """
+    feature_set_response = feature_manager_client.get_feature_set_by_name(
+        feature_set_name
+    )
+    if not feature_set_response:
+        raise FrogmlException(f"Featureset '{feature_set_name}' does not exist")
+
+    featureset_type: ProtoFeatureSetType = get_feature_type(
+        feature_set_response.feature_set.feature_set_definition.feature_set_spec.feature_set_type
+    )
+    featureset_version = _get_featureset_version(featureset_type)
+
+    return FeatureSetInfo(
+        entity_spec=feature_set_response.feature_set.feature_set_definition.feature_set_spec.entity.entity_spec,
+        feature_set_type=featureset_type,
+        feature_version=featureset_version,
+    )
+
+
+def get_feature_type(feature_set_type: ProtoFeatureSetType):
+    return getattr(feature_set_type, feature_set_type.WhichOneof("set_type"))
+
+
+def _get_featureset_version(feature_set_type: ProtoFeatureSetType) -> int:
+    """
+    Get Featureset version. Return 0 if none version is set
+    """
+    if hasattr(feature_set_type, "qwak_internal_protocol_version"):
+        return feature_set_type.qwak_internal_protocol_version
+    return 0
+
+
+def get_typed_feature(
+    feature: FeatureStoreInput,
+    feature_type: ProtoFeatureSetType,
+    featureset_version: int,
+) -> BaseFeature:
+    """
+    convert InputFeature to the relevant type
+    Args:
+        feature: Input feature to cast to the correct type
+        feature_type: the feature type as it registered
+        featureset_version: the version of the featureset
+
+    Return:
+        BaseFeature with the correct type
+    """
+    if isinstance(feature_type, ProtoBatchFeatureSet):
+        return BatchFeature(name=feature.name, entity=feature.entity)
+    elif isinstance(feature_type, ProtoBatchFeatureSetV1) and featureset_version == 0:
+        return BatchFeature(name=feature.name, entity=feature.entity)
+    elif isinstance(feature_type, ProtoBatchFeatureSetV1) and featureset_version != 0:
+        return BatchFeatureV1(name=feature.name, entity=feature.entity)
+    elif (
+        isinstance(feature_type, ProtoStreamingFeatureSetV1) and featureset_version == 0
+    ):
+        return StreamingFeature(name=feature.name, entity=feature.entity)
+    elif (
+        isinstance(feature_type, ProtoStreamingFeatureSetV1) and featureset_version != 0
+    ):
+        return StreamingFeatureV1(name=feature.name, entity=feature.entity)
+    elif isinstance(feature_type, ProtoStreamingAggregationFeatureSet):
+        return StreamingAggregationFeature(name=feature.name, entity=feature.entity)
+    else:
+        raise ValueError(
+            f"Feature set type {feature_type} with protocol version {featureset_version} is not supported for extraction"
+        )
+
+
+def get_env_to_featuresets_mapping(
+    feature_manager_client: FeatureRegistryClient,
+) -> Dict[str, Dict[str, FeatureSetInfo]]:
+    """
+    Get mapping of the account's environment name to feature sets and converts each FeatureSet to FeatureSetInfo
+    :param feature_manager_client:
+    :return: Dict[str, List[FeatureSetInfo]] where the key is the environment name and the value is a list of FeatureSetInfo
+    """
+    env_to_featuresets_mapping = feature_manager_client.get_env_to_featuresets_mapping()
+
+    # Convert the mapping to a dictionary of type Dict[str, List[FeatureSetInfo]]
+    result = {}
+    for (
+        env_name,
+        feature_sets_infos,
+    ) in env_to_featuresets_mapping.env_to_feature_set_mapping.items():
+        env_name = env_name.lower()
+        result[env_name] = {}
+        for feature_set in feature_sets_infos.feature_sets:
+            featureset_type: ProtoFeatureSetType = get_feature_type(
+                feature_set.feature_set_definition.feature_set_spec.feature_set_type
+            )
+            featureset_version = _get_featureset_version(featureset_type)
+
+            feature_set_info = FeatureSetInfo(
+                entity_spec=feature_set.feature_set_definition.feature_set_spec.entity.entity_spec,
+                feature_set_type=featureset_type,
+                feature_version=featureset_version,
+            )
+
+            feature_set_name = (
+                feature_set.feature_set_definition.feature_set_spec.name.lower()
+            )
+            result[env_name][feature_set_name] = feature_set_info
+    return result
+
+
+def get_entity_type(value_type: ProtoEntitySpec.ValueType):
+    """
+    Normalize entity by the enum
+    """
+
+    if value_type == ProtoEntitySpec.ValueType.STRING:
+        return str
+    elif value_type == ProtoEntitySpec.ValueType.INT:
+        return int
+
+
+def generate_key_unique_name(featureset_name: str):
+    """
+    Generate a name for the feature set key to be registered or treated as a Frogml entity based on its' featuresets'
+    name
+    """
+    return f"{_INTERNAL_KEY_PREFIX}_{featureset_name}_{round(time.time() * 1000)}"
+
+
+def get_batch_source_for_featureset(
+    batch_ds_name: str, feature_registry: FeatureRegistryClient
+):
+    batch_ds = feature_registry.get_data_source_by_name(batch_ds_name)
+    if not batch_ds:
+        raise FrogmlException(
+            f"Trying to register a feature set with a non registered data source -: {batch_ds_name}"
+        )
+
+    ds_spec = batch_ds.data_source.data_source_definition.data_source_spec
+
+    if ds_spec.WhichOneof("type") != "batch_source":
+        raise ValueError(
+            f"Can only register streaming backfill/batch featuresets with batch sources. "
+            f"Source {batch_ds_name} is of type {ds_spec.WhichOneof('type')}"
+        )
+
+    return ds_spec.batch_source

frogml_core/feature_store/_common/featureset_asterisk_handler.py

@@ -0,0 +1,115 @@
+import dataclasses
+import itertools
+import re
+from typing import Callable, Dict, List
+
+from frogml_core.feature_store.offline import OfflineClient
+from frogml_core.model.schema_entities import BaseFeature, Entity, FeatureStoreInput
+
+
+def is_asterisk(feature: BaseFeature) -> bool:
+    """
+    Checks whether the given Feature is a FeatureStoreInput with an '*'
+    Args:
+        feature: input feature
+    Return
+    """
+    p = "^[^.]+.\\*$"
+    # if this is not a FeatureStoreInput, no need to check
+    if isinstance(feature, FeatureStoreInput):
+        return bool(re.match(p, feature.name))
+    return False
+
+
+def inflate_feature(feature: BaseFeature, offline_fs: OfflineClient) -> List:
+    """
+    Inflates a Feature into a list. if the feature contains an '*', a list
+        a feature for each feature in the associate FeatureSet is returned. else -
+        a list containing the original feature is returned.
+    Args:
+        feature: input feature
+        offline_fs: OfflineFeatureStore to query
+    Return:
+        if an '*' exists, returns a List of FeatureStoreInuput for each feature in the feature set
+        else returns a list containing the feature.
+
+    """
+    if is_asterisk(feature):
+        featureset_name = feature.name.lower().split(".")[0]
+
+        p = f"^{featureset_name}+.*$"
+        feature_names = [
+            feature.lower()
+            for feature in offline_fs.get_columns_by_feature_set(featureset_name)
+            if re.match(p, feature.lower())
+        ]
+        return [dataclasses.replace(feature, name=name) for name in feature_names]
+
+    return [feature]
+
+
+def unpack_asterisk_features(
+    features: List[BaseFeature],
+    feature_store_generator: Callable[[], OfflineClient] = lambda: OfflineClient(),
+) -> List[BaseFeature]:
+    """
+    Handles features with an '*'.
+    If a feature of type FeatureStoreInput that matches '<featureset_name>.*' exists,
+    transforms this feature into the list of features actually present in the FeatureSet.
+    Other features remain unchanged.
+    the original order is maintained, and the internal order (inside a FeatureStoreInput
+     that has an '*') is inferred from the order of columns returned from
+     OfflineFeatureStore.get_columns_by_feature_set(...)
+     Args:
+         features: List of features
+         feature_store_generator: A function that generates an instance of OfflineFeatureStore
+                                  If no feature with an '*' is found, the function is not called.
+
+    Returns:
+        the unpacked features
+    """
+    asterisk_found = any([is_asterisk(f) for f in features])
+    if asterisk_found:
+        # at least 1 feature contains an asterisk, map each feature to
+        # a list of features (features with no '*' are mapped to a list of size 1), then flatten it.
+        offline_feature_store = feature_store_generator()
+        inflated_features = [
+            inflate_feature(feature, offline_feature_store) for feature in features
+        ]
+        return list(itertools.chain.from_iterable(inflated_features))
+
+    return features
+
+
+def unpack_asterisk_features_from_key_mapping(
+    key_to_features: Dict[str, List[str]],
+    feature_store_generator: Callable[[], OfflineClient] = lambda: OfflineClient(),
+) -> dict:
+    """
+    Handles features with an '*'.
+    If a feature name that matches '<featureset_name>.*' exists in one of the lists,
+    transforms this feature into the list of feature names actually present in the FeatureSet.
+    Other features remain unchanged.
+    the original order is maintained, and the internal order (inside a feature
+     that has an '*') is inferred from the order of columns returned from
+     OfflineFeatureStore.get_columns_by_feature_set(...)
+     Args:
+         key_to_features: dictionary of entity keys to requested features (same as in the offline store)
+             >>> key_to_features = {'uuid': ['user_purchases.*',
+             >>>                             'user_purchases.avg_purchase_amount']}
+         feature_store_generator: A function that generates an instance of OfflineFeatureStore
+                                  If no feature with an '*' is found, the function is not called.
+    Returns:
+        the unpacked features
+    """
+    inflated_key_to_features = {}
+
+    for entity, features in key_to_features.items():
+        feature_inputs = [
+            FeatureStoreInput(name=feature, entity=Entity(name=entity, type=str))
+            for feature in features
+        ]
+        unpacked = unpack_asterisk_features(feature_inputs, feature_store_generator)
+        feature_list = [feature.name for feature in unpacked]
+        inflated_key_to_features[entity] = feature_list
+    return inflated_key_to_features

frogml_core/feature_store/_common/packaging.py

@@ -0,0 +1,244 @@
+import contextlib
+import os
+import shutil
+import tempfile
+import zipfile
+from io import BytesIO
+from pathlib import Path
+from shutil import make_archive
+from typing import TYPE_CHECKING, Callable, List, Union
+from uuid import uuid4
+
+import requests
+
+from frogml_core.clients.feature_store.management_client import FeatureRegistryClient
+from frogml_core.exceptions import FrogmlException
+from frogml_core.feature_store._common.value import (
+    UPDATE_FROGML_SDK_WITH_FEATURE_STORE_EXTRA_MSG,
+)
+
+ZIP_FUNCTION_CONTENT_TYPE = "text/plain"
+
+if TYPE_CHECKING:
+    from frogml_core.feature_store.data_sources.streaming.kafka.deserialization import (
+        Deserializer,
+    )
+    from frogml_core.feature_store.feature_sets.transformations.transformations import (
+        BaseTransformation,
+    )
+
+
+def zip_source_code_dir(base_path: Path) -> Path:
+    ignore_patterns = ["*.pyc", "__pycache__"]
+    return zip_tree_to_temp_dir(base_path=base_path, ignore_patterns=ignore_patterns)
+
+
+def zip_tree_to_temp_dir(base_path: Path, ignore_patterns=None) -> Path:
+    if ignore_patterns is None:
+        ignore_patterns = []
+
+    full_path_to_archive = None
+    parent_base_path = base_path.parent if base_path.is_file() else base_path
+    tempdir = tempfile.mkdtemp()
+    dst: Path = Path(os.path.join(tempdir, parent_base_path.name))
+    try:
+        shutil.copytree(
+            src=parent_base_path,
+            dst=dst,
+            ignore=shutil.ignore_patterns(*ignore_patterns),
+        )
+
+        full_path_to_archive: Path = zip_path_dir_recursively(
+            p=dst, dest=tempfile.mkdtemp()
+        )
+        return full_path_to_archive
+    except Exception as e:
+        raise FrogmlException("Failed to copy and zip directory") from e
+    finally:
+        shutil.rmtree(tempdir, ignore_errors=True)
+        if full_path_to_archive:
+            shutil.rmtree(full_path_to_archive, ignore_errors=True)
+
+
+def zip_path_dir_recursively(p: Path, dest: Path) -> Path:
+    root_dir: Path = p.parent if p.is_file() else p
+    full_path_to_archive: str = make_archive(
+        base_name=os.path.join(dest, str(uuid4())),
+        format="zip",
+        root_dir=root_dir.parent,
+        base_dir=root_dir.name,
+    )
+    return Path(full_path_to_archive)
+
+
+def zip_function_with_base_folder(
+    functions: List[Callable], fs_object_module_dir: Path = None
+) -> bytes:
+    """
+    Prepare zip bytes of the UDF functions and it's parent folder
+    In case the functions param is a Callable it will write the pkled code to "function.pkl" (original flow)
+    If a list of functions it will write each function to {function_name}.pkl
+    Args:
+        functions: udf functions
+        fs_object_module_dir: parent directory path of the registered feature store object
+    Returns:
+        bytes representing the zip bytes of the udf functions pkl and module
+    """
+    try:
+        import cloudpickle
+    except Exception:
+        raise FrogmlException(
+            f"Missing required 'cloudpickle' dependency. {UPDATE_FROGML_SDK_WITH_FEATURE_STORE_EXTRA_MSG}"
+        )
+
+    in_memory = BytesIO()
+    with chdir(fs_object_module_dir):
+        zf = zipfile.ZipFile(in_memory, mode="w")
+        with zf:
+            if fs_object_module_dir:
+                for path in Path.cwd().glob("**/*"):
+                    str_path = (
+                        f"./{path.relative_to(Path.cwd())}"
+                        if path.is_file()
+                        else str(path.resolve())
+                    )
+                    zf.write(str_path, "code" + "/" + str_path)
+            for function in functions:
+                zf.writestr(f"{function.__name__}.pkl", cloudpickle.dumps(function))
+        in_memory.seek(0)
+    return in_memory.read()
+
+
+@contextlib.contextmanager
+def chdir(dirname: Path = None):
+    """
+    chdir to given path and return to previous cwd when done
+    Args:
+        dirname: dir to chdir into
+    """
+    current_dir = Path.cwd()
+    try:
+        if dirname is not None:
+            os.chdir(dirname)
+        yield
+    finally:
+        os.chdir(current_dir)
+
+
+def get_featureset_presign_url(
+    features_manager_client: FeatureRegistryClient,
+    feature_store_object_name: str,
+    feature_store_object_name_suffix: str,
+    transformation,
+) -> str:
+    return features_manager_client.get_presigned_url_v1(
+        feature_set_name=feature_store_object_name + feature_store_object_name_suffix,
+        transformation=transformation._to_proto(),
+    )
+
+
+def get_data_source_presign_url_v1(
+    features_manager_client: FeatureRegistryClient,
+    data_source_object_name: str,
+    data_source_name_suffix: str,
+) -> str:
+    return features_manager_client.get_data_source_presigned_url_v1(
+        data_source_name=data_source_object_name, object_name=data_source_name_suffix
+    )
+
+
+def upload_artifact(
+    feature_store_object_name: str,
+    feature_store_object_name_suffix: str,
+    functions: List[Callable],
+    feature_module_dir: Path,
+    features_manager_client: FeatureRegistryClient,
+    artifact_object: Union["BaseTransformation", "Deserializer"],
+) -> str:
+    """
+    This function is used for uploading the function code in case a UDF type transformation was used (Pandas/Koalas/Pyspark)
+    If a transformation was supplied, meaning we're running for at least batch, will use the SaaS adapted
+    function to get the presign URL.
+    """
+    uploaded_artifact_url: str
+
+    # this is done here to avoid cyclic imports
+    from frogml_core.feature_store.data_sources.streaming.kafka.deserialization import (
+        Deserializer,
+    )
+    from frogml_core.feature_store.feature_sets.transformations.transformations import (
+        BaseTransformation,
+    )
+
+    if isinstance(artifact_object, BaseTransformation):
+        presign_url = get_featureset_presign_url(
+            features_manager_client=features_manager_client,
+            feature_store_object_name=feature_store_object_name,
+            feature_store_object_name_suffix=feature_store_object_name_suffix,
+            transformation=artifact_object,
+        )
+    elif isinstance(artifact_object, Deserializer):
+        presign_url = get_data_source_presign_url_v1(
+            features_manager_client=features_manager_client,
+            data_source_object_name=feature_store_object_name,
+            data_source_name_suffix=feature_store_object_name_suffix,
+        )
+    else:
+        raise FrogmlException(
+            f"Got an unsupported artifact type: {type(artifact_object)}"
+        )
+
+    uploaded_artifact_url = base_upload_function_code(
+        presign_url, functions, feature_module_dir
+    )
+
+    return uploaded_artifact_url
+
+
+def put_presigned_url(presign_url: str, content: bytes, content_type: str) -> None:
+    http_response = requests.put(
+        presign_url,
+        headers={"content-type": content_type},
+        data=content,
+        timeout=600,
+    )
+    if 400 <= http_response.status_code < 600:
+        raise FrogmlException(
+            f"Error while uploading user function to s3. status code: {http_response.status_code}"
+        )
+
+
+def upload_to_s3(presign_url: str, in_memory_zip: bytes, content_type: str) -> str:
+    put_presigned_url(
+        presign_url=presign_url, content=in_memory_zip, content_type=content_type
+    )
+    from urllib.parse import urlparse
+
+    parsed_url = urlparse(presign_url)
+    postfix = parsed_url.path
+    bucket = parsed_url.hostname.split(".s3.")[0]
+    s3_path = f"s3://{bucket}{postfix}"  # noqa: E231
+    return s3_path
+
+
+def base_upload_function_code(
+    presign_url: str,
+    functions: List[Callable],
+    feature_module_dir: Path,
+) -> str:
+    """
+    get presign url and prepare a zip file with the functions code
+    and it's module and upload it to clients' s3 bucket
+    Args:
+        presign_url: presign url
+        functions: udf functions
+        feature_module_dir: parent directory path of the feature
+    Returns:
+        Full path to the zip file uploaded
+    """
+    in_memory_zip = zip_function_with_base_folder(functions, feature_module_dir)
+    return upload_to_s3(
+        presign_url=presign_url,
+        in_memory_zip=in_memory_zip,
+        content_type=ZIP_FUNCTION_CONTENT_TYPE,
+    )

frogml_core/feature_store/_common/source_code_spec.py

@@ -0,0 +1,62 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Optional
+
+from frogml_proto.qwak.feature_store.v1.common.source_code.source_code_pb2 import (
+    SourceCodeArtifact as ProtoSourceCodeArtifact,
+)
+from frogml_proto.qwak.feature_store.v1.common.source_code.source_code_pb2 import (
+    SourceCodeSpec as ProtoSourceCodeSpec,
+)
+from frogml_proto.qwak.feature_store.v1.common.source_code.source_code_pb2 import (
+    ZipArtifact as ProtoZipArtifact,
+)
+from frogml_core.exceptions import FrogmlException
+
+
+class SourceCodeArtifact(ABC):
+    @abstractmethod
+    def _to_proto(self) -> ProtoSourceCodeArtifact:
+        pass
+
+
+@dataclass
+class ZipArtifact(SourceCodeArtifact):
+    artifact_path: str
+    main_file: str
+
+    @classmethod
+    def _from_proto(cls, proto: ProtoSourceCodeArtifact) -> "ZipArtifact":
+        artifact_type: str = proto.WhichOneof("type")
+        if artifact_type != "zip_artifact":
+            raise FrogmlException(f"Instead of `zip_artifact` got: {artifact_type}")
+
+        return cls(
+            artifact_path=proto.zip_artifact.path,
+            main_file=proto.zip_artifact.main_file,
+        )
+
+    def _to_proto(self) -> ProtoSourceCodeArtifact:
+        return ProtoSourceCodeArtifact(
+            zip_artifact=ProtoZipArtifact(
+                path=self.artifact_path, main_file=self.main_file
+            )
+        )
+
+
+@dataclass
+class SourceCodeSpec:
+    artifact: Optional[SourceCodeArtifact] = field(default=None)
+
+    @classmethod
+    def _from_proto(cls, proto: ProtoSourceCodeSpec) -> "SourceCodeSpec":
+        artifact_type: str = proto.artifact.WhichOneof("type")
+        if artifact_type == "zip_artifact":
+            return cls(artifact=ZipArtifact._from_proto(proto.artifact))
+
+        return cls(artifact=None)
+
+    def _to_proto(self) -> ProtoSourceCodeSpec:
+        if self.artifact:
+            return ProtoSourceCodeSpec(artifact=self.artifact._to_proto())
+        return ProtoSourceCodeSpec()