Skip to content

Commit 6dc91c0

Browse files
AndrewNepogodaAndrewNepogoda
committed
Generic version of CLI
1 parent 03aadae commit 6dc91c0

File tree

7 files changed

+252
-3
lines changed

7 files changed

+252
-3
lines changed

‎cli/__init__.py

+16
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
# Copyright 2023 RelationalAI, Inc.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
__version_info__ = (0, 0, 1)
16+
__version__ = ".".join(map(str, __version_info__))

‎cli/args.py

+142
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,142 @@
1+
from argparse import ArgumentParser, Namespace, BooleanOptionalAction
2+
3+
import workflow.executor
4+
5+
6+
def parse() -> Namespace:
7+
parser = ArgumentParser()
8+
parser.add_argument(
9+
"--batch-config",
10+
help="Relative path to batch configuration json",
11+
required=True,
12+
type=str
13+
)
14+
parser.add_argument(
15+
"--batch-config-name",
16+
help="Name of to batch configuration json",
17+
required=False,
18+
type=str,
19+
default="default"
20+
)
21+
parser.add_argument(
22+
"--database",
23+
help="RAI database",
24+
required=True,
25+
type=str
26+
)
27+
parser.add_argument(
28+
"--engine",
29+
help="RAI engine",
30+
required=True,
31+
type=str
32+
)
33+
parser.add_argument(
34+
"--run-mode",
35+
help="Type of run mode",
36+
required=True,
37+
type=workflow.executor.WorkflowRunMode,
38+
choices=list(workflow.executor.WorkflowRunMode),
39+
)
40+
parser.add_argument(
41+
"--start-date", help="Start date for model data. Format: 'YYYYmmdd'",
42+
required=False,
43+
type=str
44+
)
45+
parser.add_argument(
46+
"--end-date", help="End date for model data. Format: 'YYYYmmdd'",
47+
required=True,
48+
type=str
49+
)
50+
parser.add_argument(
51+
"--dev-data-dir", help="Directory containing dev data",
52+
required=False,
53+
default="../data",
54+
type=str
55+
)
56+
parser.add_argument(
57+
"--rel-config-dir", help="Directory containing rel config files to install",
58+
required=False,
59+
default="../rel",
60+
type=str
61+
)
62+
parser.add_argument(
63+
"--env-config",
64+
help="Relative path to toml file containing environment specific RAI settings",
65+
required=False,
66+
default="../config/loader.toml",
67+
type=str
68+
)
69+
parser.add_argument(
70+
"--collapse-partitions-on-load",
71+
help="When loading each multi-part source, load all partitions (and shards) in one transaction",
72+
required=False,
73+
default=True,
74+
type=bool
75+
)
76+
parser.add_argument(
77+
"--output-root",
78+
help="Output folder path for dev mode",
79+
required=False,
80+
default="../../output",
81+
type=str
82+
)
83+
parser.add_argument(
84+
"--log-level",
85+
help="Set log level",
86+
required=False,
87+
default="INFO",
88+
type=str
89+
)
90+
parser.add_argument(
91+
"--drop-db",
92+
help="Drop RAI database before run, or not",
93+
required=False,
94+
default=True,
95+
type=bool
96+
)
97+
parser.add_argument(
98+
"--engine-size",
99+
help="Size of RAI engine",
100+
required=False,
101+
default="XS",
102+
type=str
103+
)
104+
parser.add_argument(
105+
"--cleanup-resources",
106+
help="Remove RAI engine and database after run or not",
107+
required=False,
108+
default=False,
109+
type=bool
110+
)
111+
parser.add_argument(
112+
"--disable-ivm",
113+
help="Disable IVM for RAI database",
114+
required=False,
115+
default=True,
116+
type=bool
117+
)
118+
parser.add_argument(
119+
"--recover",
120+
help="Recover a batch run starting from a FAILED step",
121+
action=BooleanOptionalAction,
122+
default=False
123+
)
124+
parser.add_argument(
125+
"--recover-step",
126+
help="Recover a batch run starting from specified step",
127+
required=False,
128+
default=False,
129+
type=str
130+
)
131+
parser.add_argument(
132+
"--rai-sdk-http-retries",
133+
help="Parameter to set http retries for rai SDK",
134+
required=False,
135+
default=3,
136+
type=int
137+
)
138+
args = parser.parse_args()
139+
# Validation
140+
if 'recover' in vars(args) and args.recover and 'recover_step' in vars(args) and args.recover_step:
141+
parser.error("`--recover` and `--recover-step` options are mutually exclusive. Use must choose only 1 option.")
142+
return args

‎cli/logger.py

+15
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
import logging
2+
3+
4+
def configure(level=logging.INFO) -> logging.Logger:
5+
# override default logging level for azure
6+
logger = logging.getLogger('azure.core')
7+
logger.setLevel(logging.ERROR)
8+
9+
logger = logging.getLogger()
10+
logger.setLevel(level)
11+
handler = logging.StreamHandler()
12+
formatter = logging.Formatter('[%(asctime)s][%(levelname)s][%(name)s] %(message)s', "%H:%M:%S")
13+
handler.setFormatter(formatter)
14+
logger.addHandler(handler)
15+
return logger.getChild("cli")

‎cli/runner.py

+71
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
1+
import time
2+
import logging
3+
import sys
4+
import tomli
5+
6+
import cli.args
7+
import cli.logger
8+
import workflow.constants
9+
import workflow.manager
10+
import workflow.common
11+
import workflow.utils
12+
import workflow.executor
13+
14+
15+
def start():
16+
# parse arguments
17+
args = cli.args.parse()
18+
# configure logger
19+
logger = cli.logger.configure(logging.getLevelName(args.log_level))
20+
try:
21+
with open(args.env_config, "rb") as fp:
22+
loader_config = tomli.load(fp)
23+
except OSError as e:
24+
logger.exception("Failed to load 'loader.toml' config.", e)
25+
sys.exit(1)
26+
# init env config
27+
env_config = workflow.common.EnvConfig.from_env_vars(loader_config)
28+
# init Workflow resource manager
29+
loader_config[workflow.constants.RAI_SDK_HTTP_RETRIES] = args.rai_sdk_http_retries
30+
resource_manager = workflow.manager.ResourceManager.init(logger, args.engine, args.database, loader_config)
31+
logger.info("Using: " + ",".join(f"{k}={v}" for k, v in vars(args).items()))
32+
try:
33+
logger.info(f"Activating batch with config from '{args.batch_config}'")
34+
start_time = time.time()
35+
# load batch config as json string
36+
batch_config_json = workflow.utils.read(args.batch_config)
37+
# create engine if it doesn't exist
38+
resource_manager.create_engine(args.engine_size)
39+
# Skip infrastructure setup during recovery
40+
if not args.recover and not args.recover_step:
41+
# Create db and disable IVM in case of enabled flag
42+
resource_manager.create_database(args.drop_db, args.disable_ivm)
43+
# Init workflow executor
44+
parameters = {
45+
workflow.constants.REL_CONFIG_DIR: args.rel_config_dir,
46+
workflow.constants.OUTPUT_ROOT: args.output_root,
47+
workflow.constants.LOCAL_DATA_DIR: args.dev_data_dir,
48+
workflow.constants.START_DATE: args.start_date,
49+
workflow.constants.END_DATE: args.end_date,
50+
workflow.constants.COLLAPSE_PARTITIONS_ON_LOAD: args.collapse_partitions_on_load
51+
}
52+
config = workflow.executor.WorkflowConfig(env_config, args.run_mode,
53+
workflow.common.BatchConfig(args.batch_config_name,
54+
batch_config_json),
55+
args.recover,
56+
args.recover_step, parameters)
57+
executor = workflow.executor.WorkflowExecutor.init(logger, config, resource_manager)
58+
end_time = time.time()
59+
executor.run()
60+
# Print execution time information
61+
executor.print_timings()
62+
logger.info(f"Infrastructure setup time is {workflow.utils.format_duration(end_time - start_time)}")
63+
64+
if args.cleanup_resources:
65+
resource_manager.cleanup_resources()
66+
except Exception as e:
67+
# Cleanup resources in case of any failure.
68+
logger.exception(e)
69+
if args.cleanup_resources:
70+
resource_manager.cleanup_resources()
71+
sys.exit(1)

‎requirements.txt

+1
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
tomli==2.0.1
12
ed25519==1.5
23
protobuf==3.20.2
34
pyarrow==6.0.1

‎setup.py

+2-1
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@
3333
],
3434
description="The RelationalAI Workflow Manager for batch runs",
3535
install_requires=[
36+
"tomli==2.0.1",
3637
"ed25519==1.5",
3738
"rai-sdk==0.6.14",
3839
"protobuf==3.20.2",
@@ -45,6 +46,6 @@
4546
long_description="Enables access to the RelationalAI REST APIs from Python",
4647
long_description_content_type="text/markdown",
4748
name="rai-workflow-manager",
48-
packages=["workflow"],
49+
packages=["workflow", "cli"],
4950
url="https://github.com/RelationalAI/rai-workflow-manager",
5051
version=workflow.__version__)

‎workflow/executor.py

+5-2
Original file line numberDiff line numberDiff line change
@@ -22,8 +22,11 @@
2222

2323

2424
class WorkflowRunMode(Enum):
25-
LOCAL = 1
26-
REMOTE = 2
25+
LOCAL = 'local'
26+
REMOTE = 'remote'
27+
28+
def __str__(self):
29+
return self.value
2730

2831

2932
class WorkflowStepState(str, Enum):

0 commit comments

Comments
 (0)