MLIA-963: Capture and handle Vela warning

Signed-off-by: Gergely Nagy <gergely.nagy@arm.com>
Change-Id: I6c8b0b74d6d35261eb0ff1a37b9577f9033be8f9

1 files changed, 12 insertions, 0 deletions

diff --git a/src/mlia/backend/vela/compiler.py b/src/mlia/backend/vela/compiler.py
index 78f97b2..b591056 100644
--- a/src/mlia/backend/vela/compiler.py
+++ b/src/mlia/backend/vela/compiler.py
@@ -6,6 +6,7 @@ from __future__ import annotations
 import logging
 import sys
 from dataclasses import dataclass
+from io import StringIO
 from pathlib import Path
 from typing import Any
 from typing import Literal
@@ -141,6 +142,9 @@ class VelaCompiler:  # pylint: disable=too-many-instance-attributes
             with redirect_output(
                 logger, stdout_level=logging.DEBUG, stderr_level=logging.DEBUG
             ):
+                tmp = sys.stdout
+                output_message = StringIO()
+                sys.stdout = output_message
                 compiler_driver(
                     nng,
                     arch,
@@ -149,8 +153,16 @@ class VelaCompiler:  # pylint: disable=too-many-instance-attributes
                     network_type,
                     output_basename,
                 )
+                sys.stdout = tmp
+                if (
+                    "Warning: SRAM target for arena memory area exceeded."
+                    in output_message.getvalue()
+                ):
+                    raise MemoryError("Model is too large and uses too much RAM")
 
             return OptimizedModel(nng, arch, compiler_options, scheduler_options)
+        except MemoryError as err:
+            raise err
         except (SystemExit, Exception) as err:
             raise RuntimeError(
                 "Model could not be optimized with Vela compiler."