Understanding the update rule for the policy in the policy iteration algorithm
A policy can be stochastic or deterministic. A deterministic policy is a function of the form $\pi_{\text{deterministic}}: S \rightarrow A$, that is, a function from the set of states to the set of actions. A stochastic policy is a map of the form $\pi_{\text{stochastic}} : S \rightarrow P(A)$, where $P(A)$ is a set of probability distributions ($P(A) = \{ p_{s_1}(A), p_{s_2}(A), \dots, p_{s_{|S|}}(A) \}$, where $p_{s_i}(A)$ is a probability distribution over the set of actions $A$ for the state $s_i$ and $|S|$ is the size of the set of states of the environment) over the set of actions $A$. A deterministic policy can be interpreted as a stochastic policy that gives the probability of $1$ to one of the available actions (and $0$ to the remaining actions), for each state.

In the case of value iteration (VI) and policy iteration (PI), the policy is deterministic, both in the policy evaluation (PE) or policy improvement (PI) steps.

In the PE step, for a specific state $s$ (inside the for loop), you use $\pi(s)$, that is, you assume that the action taken in the specific state $s$ is the greedy action, because of the $\text{argmax}$ in the PI, according to the current policy. However, in general, $\pi_k(s_i) \neq \pi_k(s_j)$, for $i \neq j$, where $k$ is the current iteration step (of PE and PI).

The update rule

$$\pi_{k+1}(s) \leftarrow arg max_a \sum_{s'}p(s' \mid s, a)[r(s, a, s') + \gamma v_k(s')]$$

computes the action that is expected to give the highest return, given the current and fixed value function $(v_k$). Have a look at the definition of expectation for discrete random variables: it is defined as a weighted sum (like in the update rule above, where the weights are $p(s' \mid s, a)$).

Value iteration is a shorter version of policy iteration. In VI, rather than performing a PI step for each state of the environment, when performing PE, the action that is expected to give the highest return is chosen. In VI, the policy is updated only once and at the end of policy evaluation. In PI, you alternate between PE and PI, and, at each PI, you update the policy.