Standalone remote client

ai_models/__main__.py

@@ -216,7 +216,7 @@ def _main(argv):
         parser.add_argument(
             "model",
             metavar="MODEL",
-            choices=available_models(),
+            choices=available_models() if "--remote" not in argv else None,
             help="The model to run",
         )
 
@@ -231,7 +231,19 @@ def _main(argv):
     args, unknownargs = parser.parse_known_args(argv)
 
     if args.models:
-        for p in sorted(available_models()):
+        if args.remote_execution:
+            from .remote import RemoteAPI
+
+            api = RemoteAPI()
+            models = api.models()
+            if len(models) == 0:
+                print(f"No remote models available on {api.url}")
+                sys.exit(0)
+            print(f"Models available on remote server {api.url}")
+        else:
+            models = available_models()
+
+        for p in sorted(models):
             print(p)
         sys.exit(0)
 
@@ -271,7 +283,12 @@ def _main(argv):
 
 
 def run(cfg: dict, model_args: list):
-    model = load_model(cfg["model"], **cfg, model_args=model_args)
+    if cfg["remote_execution"]:
+        from .remote import RemoteModel
+
+        model = RemoteModel(**cfg, model_args=model_args)
+    else:
+        model = load_model(cfg["model"], **cfg, model_args=model_args)
 
     if cfg["fields"]:
         model.print_fields()
@@ -289,10 +306,7 @@ def run(cfg: dict, model_args: list):
         sys.exit(0)
 
     try:
-        if cfg["remote_execution"]:
-            model.remote(cfg, model_args)
-        else:
-            model.run()
+        model.run()
     except FileNotFoundError as e:
         LOG.exception(e)
         LOG.error(

ai_models/model.py

@@ -10,7 +10,6 @@
 import logging
 import os
 import sys
-import tempfile
 import time
 from collections import defaultdict
 from functools import cached_property
@@ -24,7 +23,6 @@
 from .checkpoint import peek
 from .inputs import get_input
 from .outputs import get_output
-from .remote import RemoteClient
 from .stepper import Stepper
 
 LOG = logging.getLogger(__name__)
@@ -458,23 +456,6 @@ def write_input_fields(
                         check=True,
                     )
 
-    def remote(self, cfg: dict, model_args: list):
-        with tempfile.TemporaryDirectory() as tmpdirname:
-            input_file = os.path.join(tmpdirname, "input.grib")
-            output_file = os.path.join(tmpdirname, "output.grib")
-            self.all_fields.save(input_file)
-
-            client = RemoteClient(
-                input_file=input_file,
-                output_file=output_file,
-            )
-
-            client.run(cfg, model_args)
-
-            ds = cml.load_source("file", output_file)
-            for field in ds:
-                self.write(None, template=field)
-
 
 def load_model(name, **kwargs):
     return available_models()[name].load()(**kwargs)

ai_models/remote/__init__.py

@@ -0,0 +1,4 @@
+from .api import RemoteAPI
+from .model import RemoteModel
+
+__all__ = ["RemoteAPI", "RemoteModel"]