TODO: Improve this. For sure try to capture HTTP errors in more verbose logs to avoid excessive log generation if the error is for instance server raising 500 status code.

I guess to simplify this, we could just pass an instance of the exception here instead of the class and message? That would additionally also allow some exceptions which can/need to be instantiated with multiple arguments.

This should probably be of type ExecutionStatus?

My dynamic pipelines PR introduces running multiple steps in different threads, which doesn't work with this I think.

Can we somehow store the thread from which the heartbeat worker was started, and then interrupt that thread instead of the main one?

Yeah that is an important change, good point. interrupt_main will not work here, we will need to change the pattern a bit. Should I work my changes from your branch?

I'm wondering whether this RBAC check is even necessary, as running all of this will take quite some time (two calls to the DB, then a request to the RBAC service).

Is there any real harm in leaving this unprotected? I guess it would allow users potential access to the status of the step, which I'm not sure really is a concern.

True, we can probably do both authenication & authorization with pipeline tokens. Will discuss with @stefannica for directions.

Let me suggest an alternative: we could limit this endpoint to only be accessed by running pipelines.

Running pipelines (the containerized environment where the steps are running actually) use something called "a workload API token" which is only valid as long as the pipeline run itself is not yet finalized. These workload API tokens are tied to a particular pipeline run (or schedule, in case of scheduled pipelines). So we can also use their scope to limit the range of targets that they can update.

Some references:

• this is the code that verifies the pipeline scoped tokens (you can see some leeway is involved): https://github.com/zenml-io/zenml/blob/main/src/zenml/zen_server/auth.py#L406-L475
• same thing for the schedule-scoped tokens: https://github.com/zenml-io/zenml/blob/main/src/zenml/zen_server/auth.py#L363-L404

A sketch of how you can use this in your endpoint:

def update_heartbeat(
    step_run_id: UUID,
    auth_context: AuthContext = Security(authorize),
) -> StepHeartbeatResponse:

    ...
    if not auth_context.access_token or not auth_context.access_token.schedule_id and not auth_context.access_token.pipeline_run_id:
        raise AuthorizationException("Not authorized")

    if auth_context.access_token.pipeline_run_id:
        # optionally, check that the step ID is part of this run ID
    else: # if auth_context.access_token.schedule_id
        # optionally, check that the step ID is part of a run ID that was scheduled with this schedule

This will no longer rely on RBAC calls, but it might still flood the database with a lot of requests, so maybe you could also implement a mini-caching system like the ones used in the previous code references, to reduce its impact.

@stefannica That was my initial idea as well, but do we use those tokens also when running pipelines with service accounts? I thought at some point we used the API key directly when running scheduled pipelines, but I might be misremembering.

I know for sure though that there is a way to generate a generic unscoped token instead of a workload token when running a pipeline (by setting some token expiration env variable), so we'll have to think about how we handle this case.

@schustmi yes, even when running pipelines with service accounts, we generate workload API tokens scoped to the pipeline run or schedule. The only case where we use a generic unscoped token is if you set the ZENML_PIPELINE_API_TOKEN_EXPIRATION env variable. But this is a very obscure case, which I don't think we need to handle separately. In that case, we can just not run any RBAC like checks on this endpoint.

The problem is that this is a client-side env variable and I'm not sure we can recognize this case in the server endpoint that receives the heartbeat requests. It will just be a generic token that is not scoped to the run, and if we allow those to call the endpoint without any checks then everyone can do so, no?

Maybe I'm getting the logic wrong here, but this somehow seems off to me:

• We're in the main thread, and in the outermost layer we have the ContextReraise context manager, that converts KeyboardInterrupt to StepHeartBeatTerminationException.
• Here, we catch the StepHeartBeatTerminationException and potentially reraise it as a KeyboardInterrupt if the heatbeat worker isn't terminated. This will anyway be converted back to a StepHeartBeatTerminationException by the ContextReraise

Shouldn't the logic rather be the following:

• No ContextReraise needed at all

try:
  heartbeat_worker.start()
  run_step()
except KeyboardInterrupt:
  if heartbeat_worker.is_terminated:
    raise StepHeartBeatTerminationException(...)
  else:
    raise
except:
  ...
finally:
  heartbeat_worker.stop()