Fix the unit tests that use Rust FlyteIDL

.github/workflows/pythonbuild.yml

@@ -32,45 +32,116 @@ jobs:
           else
             echo "python_versions=[\"3.9\", \"3.12\"]" >> $GITHUB_ENV
           fi
-
-  # build:
-  #   needs:
-  #     - detect-python-versions
-  #   runs-on: ${{ matrix.os }}
-  #   strategy:
-  #     fail-fast: false
-  #     matrix:
-  #       os: [ubuntu-latest, windows-latest, macos-latest]
-  #       python-version: ${{fromJson(needs.detect-python-versions.outputs.python-versions)}}
-  #   steps:
-  #     - uses: actions/checkout@v4
-  #     - name: 'Clear action cache'
-  #       uses: ./.github/actions/clear-action-cache
-  #     - name: Set up Python ${{ matrix.python-version }}
-  #       uses: actions/setup-python@v4
-  #       with:
-  #         python-version: ${{ matrix.python-version }}
-  #     - name: Cache pip
-  #       uses: actions/cache@v3
-  #       with:
-  #         # This path is specific to Ubuntu
-  #         path: ~/.cache/pip
-  #         # Look to see if there is a cache hit for the corresponding requirements files
-  #         key: ${{ format('{0}-pip-{1}', runner.os, hashFiles('dev-requirements.in', 'requirements.in')) }}
-  #     - name: Install dependencies
-  #       run: |
-  #         pip install uv
-  #         make setup-global-uv
-  #         uv pip uninstall --system pandas pyarrow
-  #         uv pip freeze
-  #     - name: Test with coverage
-  #       run: |
-  #         make unit_test_codecov
-  #     - name: Codecov
-  #       uses: codecov/[email protected]
-  #       with:
-  #         fail_ci_if_error: false
-  #         files: coverage.xml
+  build:
+    needs:
+      - detect-python-versions
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [ ubuntu-latest, windows-latest, macos-latest ]
+        python-version: ${{fromJson(needs.detect-python-versions.outputs.python-versions)}}
+    steps:
+      - uses: actions/checkout@v4
+      - name: 'Clear action cache'
+        uses: ./.github/actions/clear-action-cache
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v4
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Cache pip
+        uses: actions/cache@v3
+        with:
+          # This path is specific to Ubuntu
+          path: ~/.cache/pip
+          # Look to see if there is a cache hit for the corresponding requirements files
+          key: ${{ format('{0}-pip-{1}', runner.os, hashFiles('dev-requirements.in', 'requirements.in')) }}
+      # # Install vcpkg for installing OpenSSL
+      # - name: Set up vcpkg
+      #   if: runner.os == 'Windows'
+      #   # The `openssl-sys` crate will automatically detect OpenSSL installations via Homebrew on macOS and vcpkg on Windows.
+      #   run: |
+      #     git clone --depth 1 https://github.com/microsoft/vcpkg.git
+      #     cd vcpkg
+      #     .\bootstrap-vcpkg.bat
+      # - name: Install OpenSSL
+      #   if: runner.os == 'Windows'
+      #   run: |
+      #     cd vcpkg
+      #     .\vcpkg install openssl
+      #     .\vcpkg integrate install
+      - name: Install dependencies
+        run: |
+          pip install uv maturin
+          make setup-global-uv
+          uv pip uninstall --system pandas pyarrow
+          uv pip freeze
+      - name: Fetch flyte code
+        uses: actions/checkout@v4
+        with:
+          repository: austin362667/flyte
+          fetch-depth: 0
+          ref: flyrs
+          path: "${{ github.workspace }}/flyte"
+      - name: Install Protobuf (Ubuntu)
+        if: runner.os == 'Linux'
+        run: |
+          sudo apt update && sudo apt install -y protobuf-compiler libprotobuf-dev libssl-dev
+      - name: Install Protobuf (macOS)
+        if: runner.os == 'macOS'
+        run: |
+          brew install protobuf openssl
+      - name: Install Protobuf (Windows)
+        if: runner.os == 'Windows'
+        run: |
+          choco install protoc strawberryperl make openssl -y
+      - name: Generate Rust Protobuf (Ubuntu)
+        if: runner.os == 'Linux'
+        working-directory: ${{ github.workspace }}/flyte/flyteidl
+        run: |
+          rm -rf ./gen/pb_rust/*
+          cargo run --bin gen_flyteidl
+      - name: Generate Rust Protobuf (macOS)
+        if: runner.os == 'macOS'
+        working-directory: ${{ github.workspace }}/flyte/flyteidl
+        run: |
+          rm -rf ./gen/pb_rust/*
+          cargo run --bin gen_flyteidl
+      - name: Generate Rust Protobuf (Windows)
+        if: runner.os == 'Windows'
+        shell: bash
+        working-directory: ${{ github.workspace }}/flyte/flyteidl
+        # The `openssl-sys` looks for the statically linked target by default.
+        # For dynamic linking, set VCPKGRS_DYNAMIC=1 in the environment. This will link to libraries built with vcpkg triplet x64-windows.
+        run: |
+          rm -rf ./gen/pb_rust/*
+          cargo run --bin gen_flyteidl
+      - name: Build wheels # Set up Python virtual environment at first, maturin needs this.
+        shell: bash
+        working-directory: ${{ github.workspace }}/flyte/flyteidl
+        run: |
+          python -m venv venv
+          if [[ "$RUNNER_OS" == "Windows" ]]; then
+            ./venv/Scripts/activate
+          else
+            source ./venv/bin/activate
+          fi
+          maturin build --release --out ./dist --interpreter python${{ matrix.python-version }} -m ./Cargo.toml
+      - name: Install built wheel
+        shell: bash
+        working-directory: ${{ github.workspace }}/flyte/flyteidl
+        run: |
+          pip install ./dist/*.whl --no-index --force-reinstall
+          pip freeze
+          python -c "import flyteidl_rust"
+      - name: Test with coverage
+        run: |
+          make unit_test_codecov
+      - name: Codecov
+        uses: codecov/[email protected]
+        with:
+          fail_ci_if_error: false
+          files: coverage.xml
 
   # build-with-extras:
   #   needs:

Makefile

@@ -68,9 +68,9 @@ unit_test_extras_codecov:
 unit_test:
 	# Skip all extra tests and run them with the necessary env var set so that a working (albeit slower)
 	# library is used to serialize/deserialize protobufs is used.
-	$(PYTEST_AND_OPTS) -m "not (serial or sandbox_test or hypothesis)" tests/flytekit/unit/ --ignore=tests/flytekit/unit/extras/ --ignore=tests/flytekit/unit/models --ignore=tests/flytekit/unit/extend ${CODECOV_OPTS}
+	$(PYTEST_AND_OPTS) -m "not (serial or sandbox_test or hypothesis or flyteidl_rust)" tests/flytekit/unit/ --ignore=tests/flytekit/unit/extras/ --ignore=tests/flytekit/unit/models --ignore=tests/flytekit/unit/extend ${CODECOV_OPTS}
 	# Run serial tests without any parallelism
-	$(PYTEST) -m "serial" tests/flytekit/unit/ --ignore=tests/flytekit/unit/extras/ --ignore=tests/flytekit/unit/models --ignore=tests/flytekit/unit/extend ${CODECOV_OPTS}
+	$(PYTEST) -m "serial and not flyteidl_rust" tests/flytekit/unit/ --ignore=tests/flytekit/unit/extras/ --ignore=tests/flytekit/unit/models --ignore=tests/flytekit/unit/extend ${CODECOV_OPTS}
 
 .PHONY: unit_test_hypothesis
 unit_test_hypothesis:

flytekit/clients/friendly.py

@@ -1027,7 +1027,5 @@ def get_download_signed_url(
         )
 
     def get_data(self, flyte_uri: str) -> flyteidl.service.GetDataResponse:
-        req = flyteidl.service.GetDataRequest(flyte_url=flyte_uri)
-
-        resp = self._dataproxy_stub.GetData(req, metadata=self._metadata)
+        resp = super(SynchronousFlyteClient, self).get_data(flyteidl.service.GetDataRequest(flyte_url=flyte_uri))
         return resp

flytekit/core/array_node.py

@@ -71,7 +71,7 @@ def __init__(
 
         self.metadata = None
         if isinstance(target, LaunchPlan):
-            if self._execution_mode != flyteidl.core.ArrayNode.FULL_STATE:
+            if self._execution_mode != flyteidl.array_node.ExecutionMode.FullState:
                 raise ValueError("Only execution version 1 is supported for LaunchPlans.")
             if metadata:
                 if isinstance(metadata, _workflow_model.NodeMetadata):

flytekit/core/artifact.py

@@ -277,8 +277,16 @@ def __rich_repr__(self):
         if self.value:
             if isinstance(self.value, flyteidl.core.LabelValue):
                 from flytekit.models import utils
+
                 if isinstance(self.value.value, flyteidl.label_value.Value.TimeValue):
-                    yield "Time Partition", str(utils.convert_to_datetime(seconds=self.value.value[0].seconds, nanos=self.value.value[0].nanos))
+                    yield (
+                        "Time Partition",
+                        str(
+                            utils.convert_to_datetime(
+                                seconds=self.value.value[0].seconds, nanos=self.value.value[0].nanos
+                            )
+                        ),
+                    )
                 elif isinstance(self.value.value, flyteidl.label_value.Value.InputBinding):
                     yield "Time Partition (bound to)", self.value.value[0].var
                 else:

flytekit/core/local_cache.py

@@ -40,7 +40,7 @@ def _calculate_cache_key(
 
     # Generate a stable representation of the underlying protobuf by passing `deterministic=True` to the
     # protobuf library.
-    hashed_inputs = LiteralMap(literal_map_overridden).to_flyte_idl().SerializeToString(deterministic=True)
+    hashed_inputs = LiteralMap(literal_map_overridden).to_flyte_idl().SerializeToString()
     # Use joblib to hash the string representation of the literal into a fixed length string
     return f"{task_name}-{cache_version}-{joblib.hash(hashed_inputs)}"
 

flytekit/core/type_engine.py

@@ -19,7 +19,6 @@
 import flyteidl_rust as flyteidl
 from dataclasses_json import DataClassJsonMixin, dataclass_json
 from google.protobuf.message import Message
-from marshmallow_enum import EnumField, LoadDumpOptions
 from mashumaro.codecs.json import JSONDecoder, JSONEncoder
 from mashumaro.mixins.json import DataClassJSONMixin
 from typing_extensions import Annotated, get_args, get_origin
@@ -654,7 +653,6 @@ def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type:
                 f"{expected_python_type} is not of type @dataclass, only Dataclasses are supported for "
                 "user defined datatypes in Flytekit"
             )
-
         json_str = flyteidl.DumpStruct(lv.scalar.generic)
 
         # The function looks up or creates a JSONDecoder specifically designed for the object's type.

flytekit/core/utils.py

@@ -131,7 +131,22 @@ def _get_container_definition(
 
 
 def _sanitize_resource_name(resource: "task_models.Resources.ResourceEntry") -> str:
-    return flyteidl.resources.ResourceName(resource.name).lower().replace("_", "-")
+    from flytekit.models import task as task_models
+
+    if resource.name == task_models.Resources.ResourceName.UNKNOWN:
+        return str(flyteidl.resources.ResourceName.Unknown).replace("ResourceName.", "").lower()
+    elif resource.name == task_models.Resources.ResourceName.CPU:
+        return str(flyteidl.resources.ResourceName.Cpu).replace("ResourceName.", "").lower()
+    elif resource.name == task_models.Resources.ResourceName.GPU:
+        return str(flyteidl.resources.ResourceName.Gpu).replace("ResourceName.", "").lower()
+    elif resource.name == task_models.Resources.ResourceName.MEMORY:
+        return str(flyteidl.resources.ResourceName.Memory).replace("ResourceName.", "").lower()
+    elif resource.name == task_models.Resources.ResourceName.EPHEMERAL_STORAGE:
+        import re
+
+        s = str(flyteidl.resources.ResourceName.EphemeralStorage).replace("ResourceName.", "")
+        return re.sub(r"(?<!^)(?=[A-Z])", "-", s).lower()
+    return ""
 
 
 def _serialize_pod_spec(

flytekit/exceptions/user.py

@@ -6,7 +6,7 @@
 from flytekit.exceptions.base import FlyteRecoverableException as _Recoverable
 
 
-class FlyteUserException(_FlyteException):
+class FlyteUserException(_FlyteUserException, _FlyteException):
     _ERROR_CODE = "USER:Unknown"
 
 
@@ -62,7 +62,7 @@ def __init__(self, path: str):
         super(FlyteDataNotFoundException, self).__init__(path, "File not found")
 
 
-class FlyteAssertion(_FlyteUserException, AssertionError):
+class FlyteAssertion(FlyteUserException, AssertionError):
     _ERROR_CODE = "USER:AssertionError"
 
 

flytekit/extras/accelerators.py

@@ -192,9 +192,9 @@ def to_flyte_idl(self) -> flyteidl.core.GpuAccelerator:
             return msg
 
         if self._partition_size is None:
-            msg.unpartitioned = True
+            msg.unpartitioned = flyteidl.gpu_accelerator.PartitionSizeValue.Unpartitioned(True)
         else:
-            msg.partition_size = self._partition_size
+            msg.partition_size_value = flyteidl.gpu_accelerator.PartitionSizeValue.PartitionSize(self._partition_size)
         return msg
 
 

flytekit/models/admin/common.py

@@ -34,7 +34,7 @@ def to_flyte_idl(self):
         """
         :rtype: flyteidl.admin.common_pb2.Sort
         """
-        return flyteidl.admin.Sort(key=self.key, direction=self.direction)
+        return flyteidl.admin.Sort(key=self.key, direction=int(self.direction))
 
     @classmethod
     def from_flyte_idl(cls, pb2_object):

flytekit/models/common.py

@@ -41,9 +41,14 @@ class FlyteIdlEntity(object, metaclass=FlyteType):
     def __eq__(self, other):
         import json
 
-        return isinstance(other, FlyteIdlEntity) and json.loads(other.to_flyte_idl().DumpToJsonString()) == json.loads(
-            self.to_flyte_idl().DumpToJsonString()
-        )
+        other_idl = other.to_flyte_idl()
+        self_idl = self.to_flyte_idl()
+        if isinstance(other_idl, str) and isinstance(self_idl, str):
+            return other_idl == self_idl
+        else:
+            return isinstance(other, FlyteIdlEntity) and json.loads(other_idl.DumpToJsonString()) == json.loads(
+                self_idl.DumpToJsonString()
+            )
 
     def __ne__(self, other):
         return not (self == other)
@@ -305,18 +310,15 @@ def to_flyte_idl(self):
         """
         _type = None
         if self.email:
-            _type = flytedidl.notification.Type.Email(self.email)
+            _type = flytedidl.notification.Type.Email(self.email.to_flyte_idl())
         elif self.pager_duty:
-            _type = flytedidl.notification.Type.PagerDuty(self.pager_duty)
+            _type = flytedidl.notification.Type.PagerDuty(self.pager_duty.to_flyte_idl())
         elif self.slack:
-            _type = flytedidl.notification.Type.Slack(self.slack)
+            _type = flytedidl.notification.Type.Slack(self.slack.to_flyte_idl())
 
-        return flytedidl.notification.Type(
+        return flytedidl.admin.Notification(
             phases=self.phases,
             type=_type,
-            # email=self.email.to_flyte_idl() if self.email else None,
-            # pager_duty=self.pager_duty.to_flyte_idl() if self.pager_duty else None,
-            # slack=self.slack.to_flyte_idl() if self.slack else None,
         )
 
     @classmethod

flytekit/models/core/condition.py

@@ -59,7 +59,7 @@ def to_flyte_idl(self):
         :rtype: flyteidl.core.condition_pb2.ComparisonExpression
         """
         return flyteidl.core.ComparisonExpression(
-            operator=self.operator,
+            operator=int(self.operator),
             left_value=self.left_value.to_flyte_idl(),
             right_value=self.right_value.to_flyte_idl(),
         )
@@ -174,11 +174,15 @@ def to_flyte_idl(self):
         """
         :rtype: flyteidl.core.condition_pb2.Operand
         """
-        return flyteidl.core.Operand(
-            primitive=self.primitive.to_flyte_idl() if self.primitive else None,
-            var=self.var if self.var else None,
-            scalar=self.scalar.to_flyte_idl() if self.scalar else None,
-        )
+        val = None
+        if self.primitive:
+            val = flyteidl.operand.Val.Primitive(self.primitive.to_flyte_idl())
+        elif self.var:
+            val = flyteidl.operand.Val.Var(self.var)
+        elif self.scalar:
+            val = flyteidl.operand.Val.Scalar(self.scalar.to_flyte_idl())
+
+        return flyteidl.core.Operand(val=val)
 
     @classmethod
     def from_flyte_idl(cls, pb2_object):
@@ -224,10 +228,12 @@ def to_flyte_idl(self):
         """
         :rtype: flyteidl.core.condition_pb2.BooleanExpression
         """
-        return flyteidl.core.BooleanExpression(
-            conjunction=self.conjunction.to_flyte_idl() if self.conjunction else None,
-            comparison=self.comparison.to_flyte_idl() if self.comparison else None,
-        )
+        expr = None
+        if self.conjunction:
+            expr = flyteidl.boolean_expression.Expr.Conjunction(self.conjunction.to_flyte_idl())
+        elif self.comparison:
+            expr = flyteidl.boolean_expression.Expr.Comparison(self.comparison.to_flyte_idl())
+        return flyteidl.core.BooleanExpression(expr=expr)
 
     @classmethod
     def from_flyte_idl(cls, pb2_object):