From 7d40d7f9b9bb8667f9d0741fd276dd77aa29a5db Mon Sep 17 00:00:00 2001
From: Sergi Valbuena Garrido <24233455+svalbuena@users.noreply.github.com>
Date: Tue, 6 Aug 2024 11:11:33 +0200
Subject: [PATCH] Fix Get trained models statistics API types (#2763)

* Fix Get trained models statistics API types

* Push schema changes

* Use EpochTime<UnitMillis>

(cherry picked from commit d1a056526b994e704e93a4b395b8bc8b8ee49460)
---
 output/schema/schema.json               | 22 +++++++++++++++++-----
 output/typescript/types.ts              |  4 ++--
 specification/ml/_types/TrainedModel.ts |  4 ++--
 3 files changed, 21 insertions(+), 9 deletions(-)

diff --git a/output/schema/schema.json b/output/schema/schema.json
index c10c5ec03c..005da1c94b 100644
--- a/output/schema/schema.json
+++ b/output/schema/schema.json
@@ -153121,10 +153121,13 @@
           "name": "nodes",
           "required": true,
           "type": {
-            "kind": "instance_of",
-            "type": {
-              "name": "TrainedModelDeploymentNodesStats",
-              "namespace": "ml._types"
+            "kind": "array_of",
+            "value": {
+              "kind": "instance_of",
+              "type": {
+                "name": "TrainedModelDeploymentNodesStats",
+                "namespace": "ml._types"
+              }
             }
           }
         },
@@ -153441,9 +153444,18 @@
           "name": "timestamp",
           "required": true,
           "type": {
+            "generics": [
+              {
+                "kind": "instance_of",
+                "type": {
+                  "name": "UnitMillis",
+                  "namespace": "_types"
+                }
+              }
+            ],
             "kind": "instance_of",
             "type": {
-              "name": "DateTime",
+              "name": "EpochTime",
               "namespace": "_types"
             }
           }
diff --git a/output/typescript/types.ts b/output/typescript/types.ts
index 28d8bf4b7f..3789b49da9 100644
--- a/output/typescript/types.ts
+++ b/output/typescript/types.ts
@@ -13958,7 +13958,7 @@ export interface MlTrainedModelDeploymentStats {
   error_count: integer
   inference_count: integer
   model_id: Id
-  nodes: MlTrainedModelDeploymentNodesStats
+  nodes: MlTrainedModelDeploymentNodesStats[]
   number_of_allocations: integer
   queue_capacity: integer
   rejected_execution_count: integer
@@ -13993,7 +13993,7 @@ export interface MlTrainedModelInferenceStats {
   failure_count: integer
   inference_count: integer
   missing_all_fields_count: integer
-  timestamp: DateTime
+  timestamp: EpochTime<UnitMillis>
 }
 
 export interface MlTrainedModelLocation {
diff --git a/specification/ml/_types/TrainedModel.ts b/specification/ml/_types/TrainedModel.ts
index c55b106b5f..a1ba27ed7a 100644
--- a/specification/ml/_types/TrainedModel.ts
+++ b/specification/ml/_types/TrainedModel.ts
@@ -75,7 +75,7 @@ export class TrainedModelDeploymentStats {
    * The deployment stats for each node that currently has the model allocated.
    * In serverless, stats are reported for a single unnamed virtual node.
    */
-  nodes: TrainedModelDeploymentNodesStats
+  nodes: TrainedModelDeploymentNodesStats[]
   /** The number of allocations requested. */
   number_of_allocations: integer
   /** The number of inference requests that can be queued before new requests are rejected. */
@@ -120,7 +120,7 @@ export class TrainedModelInferenceStats {
   /** The number of inference calls where all the training features for the model were missing. */
   missing_all_fields_count: integer
   /** The time when the statistics were last updated. */
-  timestamp: DateTime
+  timestamp: EpochTime<UnitMillis>
 }
 
 export class TrainedModelSizeStats {