Configuring Polaris

There's a bit of a mismatch between Spark and Polaris terminology as illustrated below:

Spark	Polaris
spark.sql.catalog.[CATALOG_NAME].warehouse	catalog
namespace	no equivalent
no equivalent	namespace

Polaris CLI

The main areas in Polaris (at least as far as the CLI is concerned) are:

A catalog has a catalog-role.

A principle has a principle-role.

The two roles are associated.

The catalog/catalog-role pairing has a number of CatalogPrivileges associated with it and these dictate what can be done (table creation, read, write etc)

Clouds

Azure

Create an App and assign it to the resource group your container ("bucket") lives in. 

Give this App's login details to Polaris as environment variables and the AppName to the Polaris catalog.

GCP

Create a service account and grant yourself the permission to use it to create tokens. Create an Application Default Credentials (ADC) JSON file for your user but this allows the app to mint access tokens as the service account.

Give Polaris the path to this ADC file as an environment variable.

(Note, I had to up the version of my Polaris to avoid this bug).

AWS

Create a role whose role policy allows users to assume it. Attach the policy allowing access to S3 to that role. 

Give your own AWS credentials to Polaris as environment variables.

KeyCloak Security

KeyCloak/OAuth

The nature of OAuth in web usage includes a frontend client like a web browser and a backend client which is the application (for example, Apache Polaris). 

Realms are isolated namespaces of security. They can be federated but are not by default.

The JWT (JSON Web Token) is associated with a realm that is passed to the backend client. It contains an ISSued field that indicates who issued it. The KeyCloak issuing it signs it with its private key.

The ISS is a URL that the embeds both the realm and the domain of the ISS. The backend client can use this URL to find a public key of KeyCloak with which it can verify the provenance of the JWT. 

This endpoint is called the Realm JWKS (JSON Web Key Set).

You can see the contents of the token with a simple Linux line like this:

echo $POLARIS_TOKEN | cut -d. -f2 | tr '_-' '/+' | base64 -d | jq

If a claim in the JWT is azp (authorized party), it identifies to whom the token was issued - typically the front end app. If it's audience, it identifies for whom the token was issued - typically the backend.

So, an example would be that the user logs in and their web browser is the azp. They ask KeyCloak for a token on behalf of the back end, the aud. KeyCloak signs both so the backend knows the token is genuinely for it and who requested it. If the azp is not in its whitelist, the backend can reject it. 

KeyCloak Configuration

To relax this constraint (say, for integration tests), Polaris can set the environment variable quarkus.oidc.token.issuer=any.

Interestingly, OAuth2 servers have a handy endpoint at .well-known/openid-configuration listing salient details. So, for KeyCloak, run:

curl -s https://YOUR_DOMAIN/realms/YOUR_REALM/.well-known/openid-configuration | jq

If your backend is in the same private network as KeyCloak, you can set KC_HOSTNAME_STRICT_BACKCHANNEL=false to make KeyCloak dynamically generate backchannel URLs based on the request's Host header - the mandatory HTTP header that defines which virtual host is the endpoint.

Note that an attacker could use Host Header Injection to trick KeyCloak into certifying a malicious domain. To stop this, the load balancer must be correctly configured to ignore inappropriate Host values plus KeyCloak should have KC_HOSTNAME_STRICT=true and KC_HOSTNAME set so even if the hacker gets past the load balancer, KeyCloak says "nice try but I know what public facing name I have". Again, you might like to relax this for testing.

Talking of the load balancer, traffic between it and KeyCloak is unencrypted so we need to tell KeyCloak communication is plain HTTP after TLS termination with KC_HTTP_ENABLED=true and set a KC_PROXY_HEADERS strategy to tell it what the load balancer is doing with the traffic. We use xforwarded so it knows our AWS ALB injects headers in a standard way.

LLMs in Java - part 1

Interpretability of LLMs is important in many real world situations. For instance, there is an EU law that says if a bank's model refuses a loan then the client has the right to know why. But neural nets are notoriously hard to interpret.

There are a few techniques:

Activation Patching: this is basically A/B testing a neural net where the activations from a clean prompt are transplanted to that of a corrupted prompt and we check if that fixes the corruption.

Logit lens: take the residual stream vector (that is, the vector that represents the working memory) and apply it to the vocabulary logits at each layer. This way, we can see where an output word (probably) starts to form.

Sparse Auto Encoders: the auto encoder is trained on the model's input and output, and deliberately made sparse for low values so prominent features become clear.

The neural net itself

Most of the Java code that iterates over the tokens lives in InferenceEngine.generateTokensGPULlama. In turn, this delegates to the Tornado graphs that are executed in TornadoVMMasterPlan.tornadoVMForwardExecuteLayered. Note that the positionHolder is an IntArray that has a single element.

However, the real heavy lifting is done in LlamaFP16FFNLayers.setupSingleFFNLayer and all the functions in TransformerComputeKernelsLayered it uses. Note, that it doesn't call them. Instead, it builds them up into a graph of lazily evaluated functions that are executed in the code mentioned above.

Note that values whose name ends with Cache are a distillation of all that has come before.

The functions that are called are (in this order):

reductionOneBlockWithLayer

reductionFinalNormalization

fusedQKVMatmulX

ropeRotationWithCacheCopy

processHeadsFlashAttention

processHeadsParallel

matrixVectorGenericWithResidual

reductionOneBlockWithLayer

reductionFinalNormalization

fusedRmsNormFFNGateUp

matrixVectorGenericWithResidual

The italicised functions configure the attention.

Let's go through them:

reductionOneBlockWithLayer calculates the root mean square of the input and takes its inverse. This is calculated per work group.

reductionFinalNormalization reduces these calculations of all the workgroups to a single value.

 

fusedQKVMatmulX multiplies the vector x by the vectors wq, wk, and wv, putting the results in q, k, and v respectively. 

ropeRotationWithCacheCopy here we see a standard, two-dimensional rotation on pairs of elements in q and k. The angle we rotate by is a function of both the word's position and which attention head this is for (which is just the global thread ID mod the head size). Note we're imposing a geometry on vectors q and k that doesn't exist before this point. 

processHeadsFlashAttention this calculates softmax(Q K^T) V. It uses tiling [previous post] as this is a massive matrix operation.

This is done for all the FFN layers plus similar calls in LogitsFP16Layer.setupLogitsTaskGraph for the logit layer. Here, the functions are (in this order):

reductionOneBlockWithLayer

reductionFinalNormalization

mapContextWithQuantizeLogits

matrixVectorGeneric

Note that these are all called for each of the layers. There appear to be 16 for Llama.

There are more layers but this is enough for my first post on the subject. More in Part 2.

Rope

Upon loading the model, RoPE.precomputeFreqsCis is called when just before the weights are loaded.

We do a 2-d rotation on pair elements of the vector, imposing a geometry that was not there before. The choice of a 2-d rotation is arbitrary but it's computationally the cheapest dimension to rotate as the number of elements in a rotation matrix is d² (where d is the number of dimensions we're working in).

Production ready Polaris

Persisting the metadata

First, create a DB. I created a Postgres RDS database in AWS and allowed it to create the VPC, subnets etc. It took me a while to work out why I could connect from my laptop but not a Polaris running in Azure: the source in the security group AWS automatically generated allowed my IP address but not Microsoft's.

You can check what IP address the greater internet sees you as with:

curl -s https://checkip.amazonaws.com

Bootstrap Polaris with:

docker run --rm -it   --env="polaris.persistence.type=relational-jdbc"   --env="quarkus.datasource.username=$DB_USERNAME"   --env="quarkus.datasource.password=$DB_PASSWORD"   --env="quarkus.datasource.jdbc.url=jdbc:postgresql://$DB_HOST:5432/polaris_db"   apache/polaris-admin-tool:latest bootstrap -r POLARIS -c POLARIS,root,s3cr3t

(You can purge the database by using the above but with arguments purge -r POLARIS)

Then you can see in Postgres:

postgres=> \c polaris_db

polaris_db=> SELECT * FROM pg_catalog.pg_tables;

     schemaname     |           tablename           | tableowner | tablespace | hasindexes | hasrules | hastriggers | rowsecurity 

--------------------+-------------------------------+------------+------------+------------+----------+-------------+-------------

 polaris_schema     | version                       | postgres   |            | t          | f        | f           | f

 polaris_schema     | entities                      | postgres   |            | t          | f        | f           | f

 polaris_schema     | grant_records                 | postgres   |            | t          | f        | f           | f

 polaris_schema     | principal_authentication_data | postgres   |            | t          | f        | f           | f

Create the database with:

CREATE DATABASE polaris_db;

CREATE USER polaris_user WITH PASSWORD 'your_secure_password';

GRANT ALL PRIVILEGES ON DATABASE polaris_db TO polaris_user;

\c polaris_db

GRANT ALL ON SCHEMA public TO polaris_user;

If you mess up your Polaris, just run:

kubectl rollout restart deployment polaris-deployment

as now the data is all in the database.

Access Control

For integration tests, you can just use the client_id and client_secret with which you set up Polaris. But if you put it in production, you'll want to create users (Principals).

"At the most basic level, Polaris' persistence layer stores Entities and Grants, where Grants define the access-control-related relationship between entities." [Apache Polaris Catalog Federation Proposal]

To access Polaris, you need a Principal. This will have Principal Roles. They need to be associated with the Catalog Roles that in turn belong to a Catalog.

REST via curl

Effective debugging of Polaris can be done by poking its REST API. There is a command line tool to do this more easily but it's a bit buggy. The latest code for me had trouble deserializing SigV4 objects.

First, you need a token

POLARIS_TOKEN=$(curl -X POST "https://$HOST/api/catalog/v1/oauth/tokens"   -H "Content-Type: application/x-www-form-urlencoded"   -d "grant_type=client_credentials&client_id=$CLIENT_ID&client_secret=$CLIENT_SECRET&scope=PRINCIPAL_ROLE:ALL" | jq -r '.access_token')

View namespaces

curl -X GET "https://$HOST/api/catalog/v1/$CATALOG_NAME/namespaces/$NAMESPACE" -H "Authorization: Bearer ${POLARIS_TOKEN}" | jq

{

  "namespace": [

    "samples"

  ],

  "properties": {

    "owner": "henryp",

    "location": "s3a://emrys-afon-bucket/samples/"

  }

}

View tables

curl -X GET "https://$HOST/api/catalog/v1/$CATALOG_NAME/namespaces/$NAMESPACE/tables"   -H "Authorization: Bearer ${POLARIS_TOKEN}"   -H "Accept: application/json"   -s | jq .

or given a table:

curl -X GET "https://$HOST/api/catalog/v1/$CATALOG_NAME/namespaces/$NAMESPACE/tables/$TABLE" -H "Authorization: Bearer ${POLARIS_TOKEN}" | jq

Clean up

Remove lingering details in the namespace with something like:

curl -X POST -H "Content-Type: application/json" -H "Authorization: Bearer ${POLARIS_TOKEN}"  "https://$HOST/api/catalog/v1/aws/namespaces/samples/properties" --data '{"removals": ["owner"] }' | jq

Delete Catalog

curl -X DELETE "https://$HOST/api/management/v1/catalogs/$CATALOG_NAME"   -H "Authorization: Bearer ${POLARIS_TOKEN}"   -H "Content-Type: application/json" -o /dev/null -s -w "%{http_code}\n"

R-squared

R² is a good way to test how well your model describes the data. It is literally

1 - (variance you predict / total variance)

It is therefore model agnostic. (For more about the maths of variance, see a previous post).

A value of 1.0 means the model describes the data perfectly, 0.0 means it is the same as guessing the mean and less than zero means it's worse than useless.

So, here's an interesting question: if my prediction for the next value in a series is the last known value, what is R²? In theory, it should be 1 and there is a mathematical argument for this. 

Some general points.

1. The formula for covariance is Σ(X_t-X̅)(Y_t-Ȳ)/N . 

Expand it and you'll see it's the same as E[XY] - E[X]E[Y].

2. We assume that a value is the sum of previous shocks mutliplied by a factor that geometrically reduces with each step. We'll call this φ, the autoregressive coefficient.

So:

y_t = ε_t + φε_t-1 + φ²ε_t-2 + ... where |φ|<1

And with a self-substitution, this becomes:

y_t = ε_t + φ y_t-1

We generally add a constant to this formula but for mathematical simplicity the rest of the argument assumes we've centred the variable on its mean - ie E[Y]=0.

3. The covariance of two unrelated distributions is 0

1. note that the joint probability distribution of x and y if they're independent is f(x,y) = f(x)f(y)

2. integrate the expected values, that is integrate xy.f(x,y) = x f(x) y f(y)

3. When you do that, you'll find E[XY] = E[X]E[Y].

4. Substituing this into the formula for covariance above, Cov(X,Y) = 0 if X and Y are independent.

4. The covariance of two distributions that are the same is the variance.

Substitute X=Y into the formula for covariance and you get Cov(X,Y)=Cov(X,X)=E[X²]-E[X]²=Var(X) as you can see in a previous blog post.

5. Take that formula in step 2, multiply both sides by y_t-1 and take the covariance of all the terms. 

Cov(ε_ty_t-1) is zero because of step 3.

Cov((y_t-1)2)=Var(Y)= E[Y²]-E[Y]² = E[Y²] because E[Y]=0 since we centred Y on its mean.

So, rearranging φ=Cov(y_t,y_t-1)/Var(y_t-1)

For a random walk:

y_t = ε_t + y_t-1

Necessarily φ=1 because the most recent step is the last plus some randomness.

So, the variance in our model is Var(ε). This is because the above formula can equivalently be written as:

y_t = y₀ + ε₁ + ε₂ + ε₃ +...

and choosing y₀=0 then:

Var(y_t) = Var(ε₁) + Var(ε₂) + Var(ε₃)+...

Var(y_t) = N Var(ε)

The total variance changes over time. Centering our walk at zero, y_t = N ε_t after N steps. So, here we find the variance is N Var(ε).

Putting these two values into the equation at the top of this post,

R² = 1 - Var(ε) / N Var(ε) = 1 - (1/N)

Naturally, as N tends to infinity, R² tends to 1.

So, a model that predicts the next value based on the last is a good one because R² is 1, right? No! This is a spurious regression. They may be useful for diagnostics (maybe) but not for models.

Linear Regression tip 

I built a decent sales model with:

MAE      :         1439240.3746               6.3954 % of mean

MAPE     :              17.0928 %

R²       :               0.9646

If you plot the residuals against the target values, you should see a random scatter plot around y=0.

This was not the case for me on my first try. There was a pattern. The residuals tended to be below zero for small target values and above zero for large. Basically, the plots were saying I was overestimating low values and underestimating large values. 

Since this was a linear regression, adding some non-linearity really helped. I both squared and loged all my regressors:

MAE      :         1293045.8398               5.1369 % of mean

MAPE     :              10.7136 %

R²       :               0.9717

This was great but the improvement didn't appear to actually come from the non-linearity. Instead, it was this line removing lots of rows:

        df = df.replace([np.inf, -np.inf], np.nan).dropna()

We need this because linear regression is intolerant to anything but a real number. Why removing the lines that had negative values (and consequently infinite loged regressors) remains a mystery. 

But it also caused another problem. As more columns were added, the probability of a row being excluded because it had a nan increased. Therefore, this little line of Python in our data pipeline had the unintended consequence of reducing the number of rows as the number of columns increased!

Polaris and Cloud tokens

Polaris rather pleasingly mints cloud tokens that are subscoped to a directory in a bucket or blob container for AWS and GCP. That is, even if the token has been hijacked, the blast radius is limited by:

• the token only allowing access to a single directory and its subfolders, not the whole bucket
• the token is no good after X minutes (where the default valu of X is 60)

There's currently an outstanding ticket to give subscoping to Azure.

The code for vending for the different clouds belongs in the implementations of PolarisStorageIntegration.getSubscopedCreds and this is where the tokens are created. You could put breakpoints in the breakpoints of:

com.google.auth.oauth2.AccessToken

com.azure.core.credential.AccessToken

software.amazon.awssdk.auth.credentials.AwsSessionCredentials 

and grab the credentials and use them on the command line (that is, entirely outside of Polaris) thus:

# AWS

AWS_ACCESS_KEY_ID=... AWS_SECRET_ACCESS_KEY=... AWS_SESSION_TOKEN=...  aws s3 ls s3://YOUR_BUCKET/DIRECTORY_FOR_TOKEN

# Azure

az storage blob list  --account-name $STORAGE_ACCOUNT   --container-name $CONTAINER --sas-token $SAS_TOKEN --prefix YOUR_DIRECTORY

# GCP

CLOUDSDK_AUTH_ACCESS_TOKEN=ya29... gcloud storage ls gs://YOUR_BUCKET/DIRECTORY_FOR_TOKEN

But even if you did, the tokens no longer work after 60 minutes and in the case of AWS and GCP, you cannot even view directories for which the token was not defined.

The state of Polaris

Huzzah! My PR for Apache Polaris has been accepted and merged with its main branch! Here are some miscellaneous notes I made as I looked at what needed to be done and how to test my code.

Federated catalogs

How Polaris handles vended credentials in federated catalogs is still an ongoing concern [Polaris mailing list]. The issue concerns who has say over what is vended. If the external catalog does not allows user X but the Polaris instance that defers to it does, is user X allowed to use that data or not?

In the ticket "Does Polaris support credential vending for external REST Catalogs?", Polaris maintainer, Alex Dutra, says:

"When the client requests credential vending, Polaris forwards the request to the remote catalog, but mints temporary credentials itself and vends them to the client. IOW, a PolarisStorageConfigurationInfo must have been configured when declaring the external catalog in Polaris, and it's this storage config that will be used for vending credentials."

Integration tests 

My GcpCatalogFederationIntegrationIT lives in runtime/service/src/cloudTest/ and unlike their counterparts in runtime/service/src/, they need to run against an already started Polaris instance (the latter start their own and run out of the box).

Note, you'll have to set:

polaris.features."ENABLE_CATALOG_FEDERATION"=true

polaris.features."ALLOW_OVERLAPPING_CATALOG_URLS"=true

in runtime/defaults/src/main/resources/application.properties and

"-Dpolaris.bootstrap.credentials=POLARIS,test-admin,test-secret",

in runtime/server/build.gradle.kts if you're running your Polaris from the source code because ServerManager hard codes the ClientCredentials.

Note that your Polaris will need to Google credentials for my code - for example:

export GOOGLE_APPLICATION_CREDENTIALS=/home/henryp/gcp.json 

Now you can run Polaris with:

./gradlew --stop && ./gradlew run