MLProfileSpec.hs

{-# LANGUAGE OverloadedStrings #-}

module Test.MLProfileSpec (spec) where

import Data.ByteString.Lazy.Char8 qualified as LBS
import Data.Map.Strict qualified as M
import Database.Bloodhound.OpenSearch3.Types
import TestsUtils.Import
import Prelude

-- | A profile response with two nodes, one carrying full per-model stats
-- and the other a bare worker_nodes routing entry. Drawn from the shape
-- documented at <https://docs.opensearch.org/latest/ml-commons-plugin/api/profile/>.
sampleProfileResponse :: LBS.ByteString
sampleProfileResponse =
  "{\
  \  \"nodes\": {\
  \    \"qTduw0FJTrmGrqMrxH0dcA\": {\
  \      \"models\": {\
  \        \"WWQI44MBbzI2oUKAvNUt\": {\"worker_nodes\": [\"KzONM8c8T4Od-NoUANQNGg\"]}\
  \      }\
  \    },\
  \    \"KzONM8c8T4Od-NoUANQNGg\": {\
  \      \"models\": {\
  \        \"WWQI44MBbzI2oUKAvNUt\": {\
  \          \"model_state\": \"DEPLOYED\",\
  \          \"predictor\": \"org.opensearch.ml.engine.algorithms.text_embedding.TextEmbeddingModel@592814c9\",\
  \          \"worker_nodes\": [\"KzONM8c8T4Od-NoUANQNGg\"],\
  \          \"predict_request_stats\": {\
  \            \"count\": 2,\
  \            \"max\": 89.978681,\
  \            \"min\": 5.402,\
  \            \"average\": 47.6903405,\
  \            \"p50\": 47.6903405,\
  \            \"p90\": 81.5210129,\
  \            \"p99\": 89.13291418999998\
  \          }\
  \        }\
  \      }\
  \    }\
  \  }\
  \}"

spec :: Spec
spec = describe "ML Commons profile API" $ do
  describe "MLProfileResponse" $ do
    it "decodes the doc sample (two nodes, one with full predict stats)" $ do
      let decoded = decode sampleProfileResponse :: Maybe MLProfileResponse
      case decoded of
        Just resp -> do
          -- Two nodes in the envelope.
          length (mlProfileResponseNodes resp) `shouldBe` 2
          -- The fully-populated node's model should report DEPLOYED and a count of 2.
          let nodes = M.elems (mlProfileResponseNodes resp)
              deployedModels =
                [ m
                | n <- nodes,
                  m <- M.elems (mlProfileNodeModels n),
                  mlProfileModelModelState m == Just ModelStateDeployed
                ]
          case deployedModels of
            (m : _) ->
              (mlProfileModelPredictRequestStats m >>= mlProfileStatsCount)
                `shouldBe` Just 2
            [] -> expectationFailure "expected a DEPLOYED model profile"
        Nothing -> expectationFailure "expected MLProfileResponse decode"

    it "decodes an empty envelope" $
      decode "{\"nodes\":{}}" `shouldBe` Just (MLProfileResponse {mlProfileResponseNodes = M.empty})

  describe "MLProfileModel" $
    it "round-trips a model with stats" $ do
      let m =
            MLProfileModel
              { mlProfileModelModelState = Just ModelStateDeployed,
                mlProfileModelPredictor = Just "cls@1",
                mlProfileModelWorkerNodes = Just ["n1"],
                mlProfileModelPredictRequestStats =
                  Just
                    MLProfilePredictRequestStats
                      { mlProfileStatsCount = Just 1,
                        mlProfileStatsMax = Just 1.0,
                        mlProfileStatsMin = Just 1.0,
                        mlProfileStatsAverage = Just 1.0,
                        mlProfileStatsP50 = Just 1.0,
                        mlProfileStatsP90 = Just 1.0,
                        mlProfileStatsP99 = Just 1.0
                      }
              }
      decode (encode m) `shouldBe` Just m

  describe "MLProfileRequest" $
    it "encodes only the supplied filters" $ do
      let req =
            MLProfileRequest
              { mlProfileRequestNodeIds = Just ["KzONM8c8T4Od-NoUANQNGg"],
                mlProfileRequestModelIds = Nothing,
                mlProfileRequestTaskIds = Nothing,
                mlProfileRequestReturnAllTasks = Just True,
                mlProfileRequestReturnAllModels = Just True
              }
      encode req
        `shouldBe` "{\"node_ids\":[\"KzONM8c8T4Od-NoUANQNGg\"],\"return_all_models\":true,\"return_all_tasks\":true}"