display sizes

saturn_client/cli/commands.py

@@ -1,5 +1,6 @@
 import json
 import logging
+import pprint
 
 from saturn_client.run import batch, setup_file_syncs, split
 
@@ -21,7 +22,7 @@
     ResourceStatus,
     SaturnConnection,
     ResourceType,
-    SaturnHTTPError,
+    SaturnHTTPError, ServerOptionTypes,
 )
 
 
@@ -476,6 +477,14 @@ def batch_cli(input_file):
     batch_info = deserialize(input_file)
     batch(batch_info)
 
+@cli.command("options")
+@click.option("--option-type", default=ServerOptionTypes.SIZES)
+@click.option("--glob")
+def options_cli(option_type: str = ServerOptionTypes.SIZES, glob: Optional[str] = None):
+    client = SaturnConnection()
+    results = client.list_options(option_type, glob=glob)
+    for r in results:
+        click.echo(pprint.pformat(r))
 
 @cli.command("split")
 @click.argument("recipe_template")

saturn_client/core.py

@@ -89,6 +89,16 @@ def lookup(cls, value: str):
         raise SaturnError(f'resource type "{value}" not found')
 
 
+class ServerOptionTypes:
+    AUTO_SHUTOFF = "auto_shutoff"
+    DISK_SPACE = "disk_space"
+    SIZES = "sizes"
+
+    @classmethod
+    def values(cls) -> List[str]:
+        return [cls.AUTO_SHUTOFF, cls.DISK_SPACE, cls.SIZES]
+
+
 class ResourceStatus:
     """
     Enum for resource statuses
@@ -206,6 +216,24 @@ def __init__(
         # test connection to raise errors early
         self._saturn_version = self._get_saturn_version()
 
+    def list_options(self, option_type: str, glob: Optional[str] = None) -> List:
+        if option_type not in ServerOptionTypes.values():
+            raise ValueError(f"unknown option {option_type}. must be one of {ServerOptionTypes.values()}")
+        url = urljoin(self.url, "api/info/servers")
+        response = requests.get(url, headers=self.settings.headers)
+        if not response.ok:
+            raise SaturnHTTPError.from_response(response)
+        results = response.json()[option_type]
+        if option_type != ServerOptionTypes.SIZES:
+            if glob:
+                results = [x for x in results if fnmatch(x, glob)]
+        else:
+            results = results.values()
+            if glob:
+                results = [x for x in results if fnmatch(x["name"], glob)]
+            results = sorted(results, key=lambda x: (x['gpu'], x['cores']))
+        return results
+
     @property
     def orgs(self) -> List[Dict[str, Any]]:
         url = urljoin(self.url, "api/orgs")