0.12.5rc1

Author: yashbonde

nbox/__init__.py

@@ -14,3 +14,4 @@
 from nbox.hyperloop.common.common_pb2 import Resource
 from nbox.nbxlib.logger import lo
 from nbox.projects import Project
+from nbox.auth import secret, AuthConfig, auth_info_pb, inside_pod

nbox/assets/exe.jinja

@@ -54,5 +54,9 @@ if __name__ == "__main__":
   nbxlet = NBXLet(op = get_op(cloud = True))
   fire.Fire({
     "run": nbxlet.run,     # NBX-Jobs
-    "serve": partial(nbxlet.serve, model_name = "{{ model_name }}"), # NBX-Deploy
+    "serve": partial(
+      nbxlet.serve,
+      model_name = "{{ model_name }}",
+      serving_type = "{{ serving_type }}"
+    ), # NBX-Deploy
   })

nbox/auth.py

@@ -18,7 +18,7 @@
 from functools import lru_cache
 
 import nbox.utils as U
-from nbox.utils import join, logger
+from nbox.utils import join, logger, lo
 
 
 class AuthConfig():
@@ -33,7 +33,7 @@ class AuthConfig():
 
   # things for the pod
   nbx_pod_run = "run"
-  # nbx_pod_deploy = "deploy"
+  nbx_pod_deploy = "deploy"
 
   def items():
     return [AuthConfig.workspace_id, AuthConfig.workspace_name, AuthConfig.cache]
@@ -208,7 +208,13 @@ def init_secret():
   # add any logic here for creating secrets
   if not U.env.NBOX_NO_AUTH(False):
     secret = NBXClient()
-    logger.info(f"Current workspace id: {secret(AuthConfig.workspace_id)} ({secret(AuthConfig.workspace_name)})")
+    logger.info(lo(
+      f"workspace details",
+      workspace_id = secret.workspace_id,
+      workspace_name =  AuthConfig.workspace_name,
+      token_present = len(secret.access_token) > 0,
+      nbx_url = secret.nbx_url,
+    ))
     return secret
   else:
     logger.info(f"Skipping authentication as NBOX_NO_AUTH is set to True")
@@ -231,4 +237,4 @@ def auth_info_pb():
   )
 
 def inside_pod():
-  return secret(AuthConfig.nbx_pod_run, False)
+  return secret(AuthConfig.nbx_pod_run, False) or secret(AuthConfig.nbx_pod_deploy, False)

nbox/cli.py

@@ -32,6 +32,7 @@
 from nbox.nbxlib.fire import NBXFire
 from nbox.projects import Project
 from nbox.utils import logger, lo
+from nbox.plugins.base import PluginCLI
 
 class Config(object):
   def update(self, workspace_id: str):
@@ -182,6 +183,7 @@ def main():
     "lmao"     : LmaoCLI,
     "login"    : login,
     "open"     : open_home,
+    "plugins"  : PluginCLI,
     "projects" : Project,
     "relics"   : Relics,
     "serve"    : Serve,

nbox/jobs.py

@@ -24,7 +24,7 @@
 from nbox.hyperloop.jobs.nbox_ws_pb2 import JobRequest
 from nbox.hyperloop.jobs.job_pb2 import Job as JobProto
 from nbox.hyperloop.jobs.dag_pb2 import DAG as DAGProto
-from nbox.hyperloop.common.common_pb2 import Resource, Code
+from nbox.hyperloop.common.common_pb2 import Resource
 from nbox.hyperloop.jobs.nbox_ws_pb2 import ListJobsRequest, ListJobsResponse, UpdateJobRequest
 from nbox.hyperloop.deploy.serve_pb2 import ServingListResponse, ServingRequest, Serving, ServingListRequest, ModelRequest, Model as ModelProto, UpdateModelRequest
 
@@ -164,7 +164,7 @@ def upload_job_folder(
   method: str,
   init_folder: str,
   id: str = "",
-  # name: str = "",
+  project_id: str = "",
   trigger: bool = False,
 
   # all the things for resources
@@ -206,6 +206,7 @@ def upload_job_folder(
     init_folder (str): folder with all the relevant files or ``file_path:fn_name`` pair so you can use it as the entrypoint.
     name (str, optional): Name of the job. Defaults to "".
     id (str, optional): ID of the job. Defaults to "".
+    project_id (str, optional): Project ID, if None uses the one from config. Defaults to "".
     trigger (bool, optional): If uploading a "job" trigger the job after uploading. Defaults to False.
     resource_cpu (str, optional): CPU resource. Defaults to "100m".
     resource_memory (str, optional): Memory resource. Defaults to "128Mi".
@@ -221,6 +222,8 @@ def upload_job_folder(
   from nbox.network import deploy_job, deploy_serving
   import nbox.nbxlib.operator_spec as ospec
   from nbox.nbxlib.serving import SupportedServingTypes as SST
+  from nbox.projects import Project
+  
   OT = ospec.OperatorType
 
   if method not in OT._valid_deployment_types():
@@ -231,6 +234,17 @@ def upload_job_folder(
     raise ValueError(f"Trigger can only be used with '{OT.JOB}' or '{OT.SERVING}'")
   if model_name and method != OT.SERVING:
     raise ValueError(f"model_name can only be used with '{OT.SERVING}'")
+  
+  # get the correct ID based on the project_id
+  if (not project_id and not id) or (project_id and id):
+    raise ValueError("Either --project-id or --id must be present")
+  if project_id:
+    p = Project(project_id)
+    if method == OT.JOB:
+      id = p.get_job_id()
+    else:
+      id = p.get_deployment_id()
+    logger.info(f"Using project_id: {project_id}, found id: {id}")
 
   if ":" not in init_folder:
     # this means we are uploading a traditonal folder that contains a `nbx_user.py` file
@@ -274,8 +288,8 @@ def upload_job_folder(
   if method == OT.SERVING:
     if serving_type not in SST.all():
       raise ValueError(f"Invalid serving_type: {serving_type}, should be one of {SST.all()}")
-    if serving_type == SST.FASTAPI:
-      logger.warning(f"You have selected serving_type='{SST.FASTAPI}', this assumes the object: {fn_name} is a FastAPI app")
+    if serving_type == SST.FASTAPI or serving_type == SST.FASTAPI_V2:
+      logger.warning(f"You have selected serving_type='{serving_type}', this assumes the object: {fn_name} is a FastAPI app")
       init_code = fn_name
       perform_tea = False
       load_operator = False
@@ -311,7 +325,8 @@ def upload_job_folder(
     "file_name": file_name,
     "fn_name": fn_name,
     "init_code": init_code,
-    "load_operator": load_operator
+    "load_operator": load_operator,
+    "serving_type": serving_type,
   }
 
   # create a requirements.txt file if it doesn't exist with the latest nbox version

nbox/lib/dist.py

@@ -6,6 +6,7 @@
 """
 
 import os
+import json
 
 import nbox.utils as U
 from nbox.utils import logger, lo
@@ -16,7 +17,7 @@
 from nbox.hyperloop.jobs.job_pb2 import Job
 from nbox.nbxlib.serving import serve_operator
 
-from nbox.lmao import ExperimentConfig, LMAO_RM_PREFIX
+from nbox.lmao import ExperimentConfig, LiveConfig, LMAO_RM_PREFIX, LMAO_SERVING_FILE
 from nbox.projects import Project, ProjectState
 
 # Manager
@@ -100,11 +101,11 @@ def run(self):
       elif run_tag.startswith(SILK_RM_PREFIX):
         # 27/03/2023: We are adding a new job type called Silk
         trace_id = run_tag[len(SILK_RM_PREFIX):]
-        logger.info(f"Running trace: {trace_id}")
         kwargs = {"trace_id": trace_id}
 
       # call the damn thing
       st = U.SimplerTimes.get_now_i64()
+      logger.debug(f"Args: {args}\nKwargs: {kwargs}")
       out = self.op(*args, **kwargs)
 
       # save the output to the relevant place, LMAO jobs are not saved to the relic
@@ -130,36 +131,23 @@ def run(self):
         tracer._rpc(f"RPC error in ending job {job_id}")
       U.hard_exit_program()
 
-  def serve(self, host: str = "0.0.0.0", port: int = 8000, *, model_name: str = None):
+  def serve(self, **serve_kwargs):
     """Run a serving API endpoint"""
-
-    run_tag = os.getenv("NBOX_RUN_METADATA", "")
-    logger.info(f"Run Tag: {run_tag}")
-
-    # Unlike a run above where it is only going to be triggered once and all the metadata is already indexed
-    # in the DB, this is not the case with deploy. But with deploy we can get away with something much simpler
-    # which is using a more complicated ModelProto.metadata object
-    # init_folder = U.env.NBOX_JOB_FOLDER("")
-    # if not init_folder:
-    #   raise RuntimeError("NBOX_JOB_FOLDER not set")
-    # if not os.path.exists(init_folder):
-    #   raise RuntimeError(f"NBOX_JOB_FOLDER {init_folder} does not exist")
-
-    # fp_bin = U.join(init_folder, "model_proto.msg")
-    # serving_type = SST.NBOX
-    # if fp_bin:
-    #   model_proto: ModelProto = read_file_to_binary(fp_bin, ModelProto())
-    #   logger.info(model_proto)
-    #   serving_type = model_proto.metadata.get("serving_type", SST.NBOX)
+    # run_tag = os.getenv("NBOX_RUN_METADATA", "")
+    # logger.info(f"Run Tag: {run_tag}")
+
+    # while we prepare to ship run_tag for serving, we can leverage the existing code to provide the exact
+    # same functionality
+    if os.path.exists(LMAO_SERVING_FILE):
+      logger.info(f"Found {LMAO_SERVING_FILE}")
+      with open(LMAO_SERVING_FILE, "r") as f:
+        cfg = LiveConfig.from_json(f.read())
+        ProjectState.project_id = cfg.get("project_id")
+        ProjectState.serving_id = cfg.get("serving_id")
+        logger.info(lo("Project data:", **ProjectState.data))
 
     try:
-      serve_operator(
-        op_or_app = self.op,
-        # serving_type = serving_type,
-        host = host,
-        port = port,
-        model_name = model_name
-      )
+      serve_operator(op_or_app = self.op, **serve_kwargs)
     except Exception as e:
       U.log_traceback()
       logger.error(f"Failed to serve operator: {e}")

nbox/lmao/__init__.py

@@ -1,6 +1,6 @@
 from nbox.lmao.exp import Lmao, _lmaoConfig
 from nbox.lmao.live import LmaoLive
 from nbox.lmao.cli import LmaoCLI
-from nbox.lmao.common import get_lmao_stub, get_git_details, get_record, ExperimentConfig, LMAO_RM_PREFIX
+from nbox.lmao.common import get_lmao_stub, get_git_details, get_record, ExperimentConfig, LMAO_RM_PREFIX, LiveConfig, LMAO_SERVING_FILE
 from nbox.lmao.proto import lmao_v2_pb2
 from nbox.lmao.lmao_rpc_client import LMAO_Stub

nbox/lmao/common.py

@@ -110,6 +110,18 @@ def get_project(project_id: str) -> ProjectProto:
   ))
   return project
 
+
+def resource_from_dict(d: Dict[str, Any]):
+  return Resource(
+    cpu = str(d["cpu"]),
+    memory = str(d["memory"]),
+    disk_size = str(d["disk_size"]),
+    gpu = str(d["gpu"]),
+    gpu_count = str(d["gpu_count"]),
+    timeout = int(d["timeout"]),
+    max_retries = int(d["max_retries"]),
+  )
+
 """
 Common structs
 """
@@ -159,24 +171,59 @@ def to_json(self):
   @classmethod
   def from_json(cls, json_str):
     d = json.loads(json_str)
-    # print(d["resource"])
-    d["resource"] = Resource(
-      cpu = str(d["resource"]["cpu"]),
-      memory = str(d["resource"]["memory"]),
-      disk_size = str(d["resource"]["disk_size"]),
-      gpu = str(d["resource"]["gpu"]),
-      gpu_count = str(d["resource"]["gpu_count"]),
-      timeout = int(d["resource"]["timeout"]),
-      max_retries = int(d["resource"]["max_retries"]),
-    )
+    d["resource"] = resource_from_dict(d["resource"])
     return cls(**d)
 
 
+
+class LiveConfig:
+  def __init__(
+    self,
+    resource: Resource,
+    cli_comm: str,
+    enable_system_monitoring: bool = False,
+  ):
+    self.resource = resource
+    self.cli_comm = cli_comm
+    self.enable_system_monitoring = enable_system_monitoring
+    self.keys = set()
+
+  def to_dict(self):
+    out = {
+      "resource": mpb.message_to_dict(self.resource),
+      "cli_comm": self.cli_comm,
+      "enable_system_monitoring": self.enable_system_monitoring,
+    }
+    for k in self.keys:
+      out[k] = getattr(self, k)
+    return out
+  
+  def to_json(self):
+    return json.dumps(self.to_dict())
+  
+  @classmethod
+  def from_json(cls, json_str) -> 'LiveConfig':
+    d = json.loads(json_str)
+    d["resource"] = resource_from_dict(d["resource"])
+    _cls = cls(**d)
+    for k in d:
+      if k not in ["resource", "cli_comm", "enable_system_monitoring"]:
+        _cls.add(k, d[k])
+  
+  def add(self, key, value):
+    setattr(self, key, value)
+    self.keys.add(key)
+
+  def get(self, key):
+    return getattr(self, key)
+
+
 """
 Constants
 """
 
 # do not change these it can become a huge pain later on
 LMAO_RELIC_NAME = "experiments"
 LMAO_RM_PREFIX = "NBXLmao-"
+LMAO_SERVING_FILE = "NBXLmaoServingCfg.json"
 LMAO_ENV_VAR_PREFIX = "NBX_LMAO_"

nbox/lmao/live.py

@@ -165,3 +165,6 @@ def log(self, y: Dict[str, Union[int, float, str]]):
       raise Exception("Server Error")
 
     self._total_logged_elements += 1
+
+
+# helper class

nbox/nbxlib/serving.py

@@ -32,14 +32,15 @@
 class SupportedServingTypes():
   NBOX = "nbox"
   FASTAPI = "fastapi"
+  FASTAPI_V2 = "fastapi_v2"  
 
   def all():
-    return [SupportedServingTypes.NBOX, SupportedServingTypes.FASTAPI]
+    return [getattr(SupportedServingTypes, x) for x in dir(SupportedServingTypes) if not x.startswith("__") and not x == "all"]
 
 
 def serve_operator(
   op_or_app: Operator,
-  # serving_type: str,
+  serving_type: str,
   host: str = "0.0.0.0",
   port: int = 8000,
   *,
@@ -88,10 +89,18 @@ async def who_are_you():
     for route, fn in get_fastapi_routes(op_or_app):
       app.add_api_route(route, fn, methods=["POST"], response_class=JSONResponse)
 
-  elif type(op_or_app) == FastAPI:
+  elif serving_type == SupportedServingTypes.FASTAPI:
     # app.mount("/x", op_or_app)
+    logger.info("Mounting FastAPI app at /x/...")
     app.add_api_route("/x", op_or_app, methods=["POST"], response_class=JSONResponse)
 
+  elif serving_type == SupportedServingTypes.FASTAPI_V2:
+    # load all the paths from op_or_app and add them to the app except / and /metadata
+    logger.info("Loading FastAPI app routes...")
+    for route in op_or_app.routes:
+      if route.path not in ["/", "/metadata"]:
+        app.add_api_route(route.path, route.endpoint, methods=route.methods, response_class=JSONResponse)
+
   # TODO: @yashbonde can we mount flask/django? https://fastapi.tiangolo.com/advanced/wsgi/
 
   else:

nbox/nbxlib/tracer.py

@@ -91,6 +91,7 @@ def init(self, run_data, start_heartbeat):
 
     # start heartbeat in a different thread
     if start_heartbeat:
+      self.thread_stop = threading.Event()
       self.thread = threading.Thread(target=self.hearbeat_thread_worker)
       self.thread.start()
 
@@ -128,7 +129,9 @@ def __call__(self, node: Node, verbose: bool = False):
   def hearbeat_thread_worker(self):
     while True:
       self._rpc()
-      sleep(self.heartbeat_every)
+      for _ in range(self.hearbeat_thread_worker):
+        # in future add a way to stop the thread
+        sleep(1)
 
   def stop(self):
     if not self.network_tracer: