Expose deployment version on contexts

Author: Sushisource

temporalio/bridge/Cargo.lock

@@ -5,7 +5,7 @@ use log::error;
 use prost::Message;
 use pyo3::exceptions::{PyException, PyRuntimeError, PyValueError};
 use pyo3::prelude::*;
-use pyo3::types::{PyBytes, PyTuple};
+use pyo3::types::{PyBytes, PyDict, PyTuple};
 use std::collections::HashMap;
 use std::collections::HashSet;
 use std::marker::PhantomData;
@@ -91,12 +91,31 @@ pub struct LegacyBuildIdBased {
 }
 
 /// Recreates [temporal_sdk_core_api::worker::WorkerDeploymentVersion]
-#[derive(FromPyObject)]
+#[derive(FromPyObject, Clone)]
 pub struct WorkerDeploymentVersion {
     pub deployment_name: String,
     pub build_id: String,
 }
 
+impl IntoPy<Py<PyAny>> for WorkerDeploymentVersion {
+    fn into_py(self, py: Python) -> Py<PyAny> {
+        let dict = PyDict::new(py);
+        dict.set_item("deployment_name", self.deployment_name)
+            .unwrap();
+        dict.set_item("build_id", self.build_id).unwrap();
+        dict.into()
+    }
+}
+
+impl From<temporal_sdk_core_api::worker::WorkerDeploymentVersion> for WorkerDeploymentVersion {
+    fn from(version: temporal_sdk_core_api::worker::WorkerDeploymentVersion) -> Self {
+        WorkerDeploymentVersion {
+            deployment_name: version.deployment_name,
+            build_id: version.build_id,
+        }
+    }
+}
+
 #[derive(FromPyObject)]
 pub struct TunerHolder {
     workflow_slot_supplier: SlotSupplier,
@@ -134,7 +153,7 @@ pub struct SlotReserveCtx {
     #[pyo3(get)]
     pub worker_identity: String,
     #[pyo3(get)]
-    pub worker_build_id: String,
+    pub worker_deployment_version: Option<WorkerDeploymentVersion>,
     #[pyo3(get)]
     pub is_sticky: bool,
 }
@@ -150,7 +169,7 @@ impl SlotReserveCtx {
             },
             task_queue: ctx.task_queue().to_string(),
             worker_identity: ctx.worker_identity().to_string(),
-            worker_build_id: ctx.worker_build_id().to_string(),
+            worker_deployment_version: ctx.worker_deployment_version().clone().map(Into::into),
             is_sticky: ctx.is_sticky(),
         }
     }

temporalio/bridge/sdk-core

@@ -8,6 +8,7 @@
 from typing_extensions import TypeAlias
 
 import temporalio.bridge.worker
+from temporalio.common import WorkerDeploymentVersion
 
 _DEFAULT_RESOURCE_ACTIVITY_MAX = 500
 
@@ -55,7 +56,7 @@ class ResourceBasedSlotConfig:
     ramp_throttle: Optional[timedelta] = None
     """Minimum time we will wait (after passing the minimum slots number) between handing out new slots in milliseconds.
     Defaults to 0 for workflows and 50ms for activities.
-    
+
     This value matters because how many resources a task will use cannot be determined ahead of time, and thus the
     system should wait to see how much resources are used before issuing more slots."""
 
@@ -100,11 +101,22 @@ class SlotReserveContext(Protocol):
     """The name of the task queue for which this reservation request is associated."""
     worker_identity: str
     """The identity of the worker that is requesting the reservation."""
-    worker_build_id: str
-    """The build id of the worker that is requesting the reservation."""
+    worker_deployment_version: Optional[WorkerDeploymentVersion]
+    """The deployment version of the worker that is requesting the reservation, if any."""
     is_sticky: bool
     """True iff this is a reservation for a sticky poll for a workflow task."""
 
+    @property
+    def worker_build_id(self) -> str:
+        """The build id of the worker that is requesting the reservation.
+
+        .. warning::
+            Deprecated, use :py:attr:`worker_deployment_version` instead.
+        """
+        if not self.worker_deployment_version:
+            return ""
+        return self.worker_deployment_version.build_id
+
 
 # WARNING: This must match Rust worker::WorkflowSlotInfo
 @runtime_checkable

temporalio/bridge/src/worker.rs

@@ -50,6 +50,7 @@
 import temporalio.api.sdk.v1
 import temporalio.bridge.proto.activity_result
 import temporalio.bridge.proto.child_workflow
+import temporalio.bridge.proto.common
 import temporalio.bridge.proto.workflow_activation
 import temporalio.bridge.proto.workflow_commands
 import temporalio.bridge.proto.workflow_completion
@@ -211,7 +212,9 @@ def __init__(self, det: WorkflowInstanceDetails) -> None:
         self._primary_task: Optional[asyncio.Task[None]] = None
         self._time_ns = 0
         self._cancel_requested = False
-        self._current_build_id = ""
+        self._deployment_version_for_current_task: Optional[
+            temporalio.bridge.proto.common.WorkerDeploymentVersion
+        ] = None
         self._current_history_length = 0
         self._current_history_size = 0
         self._continue_as_new_suggested = False
@@ -351,7 +354,9 @@ def activate(
             else temporalio.api.enums.v1.VersioningBehavior.VERSIONING_BEHAVIOR_UNSPECIFIED
         )
         self._current_activation_error: Optional[Exception] = None
-        self._current_build_id = act.build_id_for_current_task
+        self._deployment_version_for_current_task = (
+            act.deployment_version_for_current_task
+        )
         self._current_history_length = act.history_length
         self._current_history_size = act.history_size_bytes
         self._continue_as_new_suggested = act.continue_as_new_suggested
@@ -985,7 +990,19 @@ def workflow_extern_functions(self) -> Mapping[str, Callable]:
         return self._extern_functions
 
     def workflow_get_current_build_id(self) -> str:
-        return self._current_build_id
+        if not self._deployment_version_for_current_task:
+            return ""
+        return self._deployment_version_for_current_task.build_id
+
+    def workflow_get_current_deployment_version(
+        self,
+    ) -> Optional[temporalio.common.WorkerDeploymentVersion]:
+        if not self._deployment_version_for_current_task:
+            return None
+        return temporalio.common.WorkerDeploymentVersion(
+            build_id=self._deployment_version_for_current_task.build_id,
+            deployment_name=self._deployment_version_for_current_task.deployment_name,
+        )
 
     def workflow_get_current_history_length(self) -> int:
         return self._current_history_length

temporalio/worker/_tuning.py

@@ -480,12 +480,28 @@ def _logger_details(self) -> Mapping[str, Any]:
     def get_current_build_id(self) -> str:
         """Get the Build ID of the worker which executed the current Workflow Task.
 
-        May be undefined if the task was completed by a worker without a Build ID. If this worker is the one executing
-        this task for the first time and has a Build ID set, then its ID will be used. This value may change over the
-        lifetime of the workflow run, but is deterministic and safe to use for branching.
+        May be undefined if the task was completed by a worker without a Build ID. If this worker is
+        the one executing this task for the first time and has a Build ID set, then its ID will be
+        used. This value may change over the lifetime of the workflow run, but is deterministic and
+        safe to use for branching.
+
+        .. deprecated::
+            Use get_current_deployment_version instead.
         """
         return _Runtime.current().workflow_get_current_build_id()
 
+    def get_current_deployment_version(
+        self,
+    ) -> Optional[temporalio.common.WorkerDeploymentVersion]:
+        """Get the deployment version of the worker which executed the current Workflow Task.
+
+        May be None if the task was completed by a worker without a deployment version or build
+        id. If this worker is the one executing this task for the first time and has a deployment
+        version set, then its ID will be used. This value may change over the lifetime of the
+        workflow run, but is deterministic and safe to use for branching.
+        """
+        return _Runtime.current().workflow_get_current_deployment_version()
+
     def get_current_history_length(self) -> int:
         """Get the current number of events in history.
 
@@ -613,6 +629,11 @@ def workflow_extern_functions(self) -> Mapping[str, Callable]: ...
     @abstractmethod
     def workflow_get_current_build_id(self) -> str: ...
 
+    @abstractmethod
+    def workflow_get_current_deployment_version(
+        self,
+    ) -> Optional[temporalio.common.WorkerDeploymentVersion]: ...
+
     @abstractmethod
     def workflow_get_current_history_length(self) -> int: ...
 

temporalio/worker/_workflow_instance.py

@@ -561,6 +561,11 @@ class DeploymentVersioningWorkflowV2Pinned:
     @workflow.run
     async def run(self):
         await workflow.wait_condition(lambda: self.finish)
+        depver = workflow.info().get_current_deployment_version()
+        assert depver
+        assert depver.build_id == "2.0"
+        # Just ensuring the rust object was converted properly and this method still works
+        workflow.logger.debug(f"Dep string: {depver.to_canonical_string()}")
         return "version-v2"
 
     @workflow.signal