humalab 0.0.5__py3-none-any.whl → 0.0.6__py3-none-any.whl

humalab/humalab.py

@@ -1,23 +1,25 @@
 from contextlib import contextmanager
+import sys
+import traceback
 
 from omegaconf import OmegaConf
 
+from humalab.constants import DEFAULT_PROJECT
 from humalab.run import Run
 from humalab.humalab_config import HumalabConfig
-from humalab.humalab_api_client import HumaLabApiClient
-from humalab.constants import EpisodeStatus
+from humalab.humalab_api_client import HumaLabApiClient, RunStatus, EpisodeStatus
 import requests
 
 import uuid
 
 from collections.abc import Generator
 
-from humalab.scenario import Scenario
+from humalab.scenarios.scenario import Scenario
 
 _cur_run: Run | None = None
 
 def _pull_scenario(client: HumaLabApiClient,
-                   project_name: str,
+                   project: str,
                    scenario: str | list | dict | None = None,
                    scenario_id: str | None = None,) -> str | list | dict | None:
     if scenario_id is not None:
@@ -28,7 +30,7 @@ def _pull_scenario(client: HumaLabApiClient,
         scenario_version = int(scenario_arr[1]) if len(scenario_arr) > 1 else None
 
         scenario_response = client.get_scenario(
-            project_name=project_name,
+            project_name=project,
             uuid=scenario_real_id, version=scenario_version)
         return scenario_response["yaml_content"]
     return scenario
@@ -41,12 +43,13 @@ def init(project: str | None = None,
          tags: list[str] | None = None,
          scenario: str | list | dict | None = None,
          scenario_id: str | None = None,
+         seed: int | None=None,
+         auto_create_scenario: bool = False,
+         # num_env: int | None = None,
+
          base_url: str | None = None,
          api_key: str | None = None,
-         seed: int | None=None,
          timeout: float | None = None,
-         # num_env: int | None = None,
-         auto_create_scenario: bool = False,
          ) -> Generator[Run, None, None]:
     """
     Initialize a new HumaLab run.
@@ -63,38 +66,32 @@ def init(project: str | None = None,
         api_key: The API key for authentication.
         seed: An optional seed for scenario randomization.
         timeout: The timeout for API requests.
-        # num_env: The number of parallel environments to run. (Not supported yet.)
         auto_create_scenario: Whether to automatically create the scenario if it does not exist.
+        # num_env: The number of parallel environments to run. (Not supported yet.)
     """
     global _cur_run
     run = None
     try:
-        humalab_config = HumalabConfig()
-        project = project or "default"
+        project = project or DEFAULT_PROJECT
         name = name or ""
         description = description or ""
         id = id or str(uuid.uuid4())
 
-        base_url = base_url or humalab_config.base_url
-        api_key = api_key or humalab_config.api_key
-        timeout = timeout or humalab_config.timeout
-
         api_client = HumaLabApiClient(base_url=base_url,
                                       api_key=api_key,
                                       timeout=timeout)
         final_scenario = _pull_scenario(client=api_client, 
-                                        project_name=project,
+                                        project=project,
                                         scenario=scenario, 
                                         scenario_id=scenario_id)
         
         project_resp = api_client.create_project(name=project)
 
         scenario_inst = Scenario()
-        scenario_inst.init(run_id=id, 
-                           scenario=final_scenario, 
+        scenario_inst.init(scenario=final_scenario, 
                            seed=seed, 
-                           episode_id=str(uuid.uuid4()),
-                           #num_env=num_env
+                           scenario_id=scenario_id,
+                           #num_env=num_env,
                            )
         if scenario_id is None and scenario is not None and auto_create_scenario:
             scenario_response = api_client.create_scenario(
@@ -108,6 +105,10 @@ def init(project: str | None = None,
             run_response = api_client.get_run(run_id=id)
             api_client.update_run(
                 run_id=run_response['run_id'],
+                name=name,
+                description=description,
+                tags=tags,
+                status=RunStatus.RUNNING,
             )
 
         except requests.HTTPError as e:
@@ -137,20 +138,35 @@ def init(project: str | None = None,
             id=run_response['run_id'],
             tags=run_response.get("tags"),
             scenario=scenario_inst,
+
+            base_url=base_url,
+            api_key=api_key,
+            timeout=timeout
         )
 
         _cur_run = run
         yield run
-    finally:
-        if run:
-            run.finish()
-        
-
-def finish(status: EpisodeStatus = EpisodeStatus.PASS,
-           quiet: bool | None = None) -> None:
+    except Exception as e:
+        if _cur_run:
+            exc_type, exc_value, exc_traceback = sys.exc_info()
+            formatted_traceback = ''.join(traceback.format_exception(exc_type, exc_value, exc_traceback))
+            finish(status=RunStatus.ERRORED,
+                   err_msg=formatted_traceback)
+        raise
+    else:
+        if _cur_run:
+            print("Finishing run...")
+            finish(status=RunStatus.FINISHED)
+
+def discard() -> None:
+    finish(status=RunStatus.CANCELED)
+
+def finish(status: RunStatus = RunStatus.FINISHED,
+           err_msg: str | None = None) -> None:
     global _cur_run
     if _cur_run:
-        _cur_run.finish(status=status, quiet=quiet)
+        _cur_run.finish(status=status, err_msg=err_msg)
+        _cur_run = None
 
 def login(api_key: str | None = None,
           relogin: bool | None = None,