omdev 0.0.0.dev212__py3-none-any.whl → 0.0.0.dev213__py3-none-any.whl

Sign up to get free protection for your applications and to get access to all the features.
omdev/scripts/ci.py CHANGED
@@ -17,6 +17,8 @@ Inputs:
17
17
  import abc
18
18
  import argparse
19
19
  import asyncio
20
+ import asyncio.base_subprocess
21
+ import asyncio.subprocess
20
22
  import collections
21
23
  import contextlib
22
24
  import dataclasses as dc
@@ -39,6 +41,7 @@ import threading
39
41
  import time
40
42
  import types
41
43
  import typing as ta
44
+ import urllib.parse
42
45
 
43
46
 
44
47
  ########################################
@@ -54,6 +57,9 @@ if sys.version_info < (3, 8):
54
57
  # shell.py
55
58
  T = ta.TypeVar('T')
56
59
 
60
+ # ../../omlish/asyncs/asyncio/timeouts.py
61
+ AwaitableT = ta.TypeVar('AwaitableT', bound=ta.Awaitable)
62
+
57
63
  # ../../omlish/lite/cached.py
58
64
  CallableT = ta.TypeVar('CallableT', bound=ta.Callable)
59
65
 
@@ -70,6 +76,7 @@ ArgparseCmdFn = ta.Callable[[], ta.Optional[int]] # ta.TypeAlias
70
76
 
71
77
  # ../../omlish/lite/contextmanagers.py
72
78
  ExitStackedT = ta.TypeVar('ExitStackedT', bound='ExitStacked')
79
+ AsyncExitStackedT = ta.TypeVar('AsyncExitStackedT', bound='AsyncExitStacked')
73
80
 
74
81
  # ../../omlish/subprocesses.py
75
82
  SubprocessChannelOption = ta.Literal['pipe', 'stdout', 'devnull'] # ta.TypeAlias
@@ -113,6 +120,19 @@ class ShellCmd:
113
120
  )
114
121
 
115
122
 
123
+ ########################################
124
+ # ../../../omlish/asyncs/asyncio/timeouts.py
125
+
126
+
127
+ def asyncio_maybe_timeout(
128
+ fut: AwaitableT,
129
+ timeout: ta.Optional[float] = None,
130
+ ) -> AwaitableT:
131
+ if timeout is not None:
132
+ fut = asyncio.wait_for(fut, timeout) # type: ignore
133
+ return fut
134
+
135
+
116
136
  ########################################
117
137
  # ../../../omlish/lite/cached.py
118
138
 
@@ -208,6 +228,17 @@ class Checks:
208
228
 
209
229
  #
210
230
 
231
+ def register_on_raise_breakpoint_if_env_var_set(self, key: str) -> None:
232
+ import os
233
+
234
+ def on_raise(exc: Exception) -> None: # noqa
235
+ if key in os.environ:
236
+ breakpoint() # noqa
237
+
238
+ self.register_on_raise(on_raise)
239
+
240
+ #
241
+
211
242
  def set_exception_factory(self, factory: CheckExceptionFactory) -> None:
212
243
  self._exception_factory = factory
213
244
 
@@ -523,6 +554,18 @@ class Checks:
523
554
 
524
555
  return v
525
556
 
557
+ def not_equal(self, v: T, o: ta.Any, msg: CheckMessage = None) -> T:
558
+ if o == v:
559
+ self._raise(
560
+ ValueError,
561
+ 'Must not be equal',
562
+ msg,
563
+ Checks._ArgsKwargs(v, o),
564
+ render_fmt='%s == %s',
565
+ )
566
+
567
+ return v
568
+
526
569
  def is_(self, v: T, o: ta.Any, msg: CheckMessage = None) -> T:
527
570
  if o is not v:
528
571
  self._raise(
@@ -1205,7 +1248,7 @@ class GithubCacheServiceV1:
1205
1248
  @dc.dataclass(frozen=True)
1206
1249
  class ReserveCacheRequest:
1207
1250
  key: str
1208
- cache_size: ta.Optional[int]
1251
+ cache_size: ta.Optional[int] = None
1209
1252
  version: ta.Optional[str] = None
1210
1253
 
1211
1254
  @dc.dataclass(frozen=True)
@@ -1713,6 +1756,33 @@ class ExitStacked:
1713
1756
  return es.enter_context(cm)
1714
1757
 
1715
1758
 
1759
+ class AsyncExitStacked:
1760
+ _exit_stack: ta.Optional[contextlib.AsyncExitStack] = None
1761
+
1762
+ async def __aenter__(self: AsyncExitStackedT) -> AsyncExitStackedT:
1763
+ check.state(self._exit_stack is None)
1764
+ es = self._exit_stack = contextlib.AsyncExitStack()
1765
+ await es.__aenter__()
1766
+ return self
1767
+
1768
+ async def __aexit__(self, exc_type, exc_val, exc_tb):
1769
+ if (es := self._exit_stack) is None:
1770
+ return None
1771
+ await self._async_exit_contexts()
1772
+ return await es.__aexit__(exc_type, exc_val, exc_tb)
1773
+
1774
+ async def _async_exit_contexts(self) -> None:
1775
+ pass
1776
+
1777
+ def _enter_context(self, cm: ta.ContextManager[T]) -> T:
1778
+ es = check.not_none(self._exit_stack)
1779
+ return es.enter_context(cm)
1780
+
1781
+ async def _enter_async_context(self, cm: ta.AsyncContextManager[T]) -> T:
1782
+ es = check.not_none(self._exit_stack)
1783
+ return await es.enter_async_context(cm)
1784
+
1785
+
1716
1786
  ##
1717
1787
 
1718
1788
 
@@ -1724,6 +1794,17 @@ def defer(fn: ta.Callable) -> ta.Generator[ta.Callable, None, None]:
1724
1794
  fn()
1725
1795
 
1726
1796
 
1797
+ @contextlib.asynccontextmanager
1798
+ async def adefer(fn: ta.Callable) -> ta.AsyncGenerator[ta.Callable, None]:
1799
+ try:
1800
+ yield fn
1801
+ finally:
1802
+ await fn()
1803
+
1804
+
1805
+ ##
1806
+
1807
+
1727
1808
  @contextlib.contextmanager
1728
1809
  def attr_setting(obj, attr, val, *, default=None): # noqa
1729
1810
  not_set = object()
@@ -2277,154 +2358,619 @@ class AbstractAsyncSubprocesses(BaseSubprocesses):
2277
2358
 
2278
2359
 
2279
2360
  ########################################
2280
- # ../compose.py
2281
- """
2282
- TODO:
2283
- - fix rmi - only when not referenced anymore
2284
- """
2361
+ # ../github/curl.py
2285
2362
 
2286
2363
 
2287
2364
  ##
2288
2365
 
2289
2366
 
2290
- def get_compose_service_dependencies(
2291
- compose_file: str,
2292
- service: str,
2293
- ) -> ta.Dict[str, str]:
2294
- compose_dct = read_yaml_file(compose_file)
2367
+ class GithubServiceCurlClient:
2368
+ def __init__(
2369
+ self,
2370
+ service_url: str,
2371
+ auth_token: ta.Optional[str] = None,
2372
+ *,
2373
+ api_version: ta.Optional[str] = None,
2374
+ ) -> None:
2375
+ super().__init__()
2295
2376
 
2296
- services = compose_dct['services']
2297
- service_dct = services[service]
2377
+ self._service_url = check.non_empty_str(service_url)
2378
+ self._auth_token = auth_token
2379
+ self._api_version = api_version
2298
2380
 
2299
- out = {}
2300
- for dep_service in service_dct.get('depends_on', []):
2301
- dep_service_dct = services[dep_service]
2302
- out[dep_service] = dep_service_dct['image']
2381
+ #
2303
2382
 
2304
- return out
2383
+ _MISSING = object()
2305
2384
 
2385
+ def build_headers(
2386
+ self,
2387
+ headers: ta.Optional[ta.Mapping[str, str]] = None,
2388
+ *,
2389
+ auth_token: ta.Any = _MISSING,
2390
+ content_type: ta.Optional[str] = None,
2391
+ ) -> ta.Dict[str, str]:
2392
+ dct = {
2393
+ 'Accept': ';'.join([
2394
+ 'application/json',
2395
+ *([f'api-version={self._api_version}'] if self._api_version else []),
2396
+ ]),
2397
+ }
2306
2398
 
2307
- ##
2399
+ if auth_token is self._MISSING:
2400
+ auth_token = self._auth_token
2401
+ if auth_token:
2402
+ dct['Authorization'] = f'Bearer {auth_token}'
2308
2403
 
2404
+ if content_type is not None:
2405
+ dct['Content-Type'] = content_type
2309
2406
 
2310
- class DockerComposeRun(ExitStacked):
2311
- @dc.dataclass(frozen=True)
2312
- class Config:
2313
- compose_file: str
2314
- service: str
2407
+ if headers:
2408
+ dct.update(headers)
2315
2409
 
2316
- image: str
2410
+ return dct
2317
2411
 
2318
- cmd: ShellCmd
2412
+ #
2319
2413
 
2320
- #
2414
+ HEADER_AUTH_TOKEN_ENV_KEY_PREFIX = '_GITHUB_SERVICE_AUTH_TOKEN' # noqa
2321
2415
 
2322
- run_options: ta.Optional[ta.Sequence[str]] = None
2416
+ @property
2417
+ def header_auth_token_env_key(self) -> str:
2418
+ return f'{self.HEADER_AUTH_TOKEN_ENV_KEY_PREFIX}_{id(self)}'
2323
2419
 
2324
- cwd: ta.Optional[str] = None
2420
+ def build_cmd(
2421
+ self,
2422
+ method: str,
2423
+ url: str,
2424
+ *,
2425
+ json_content: bool = False,
2426
+ content_type: ta.Optional[str] = None,
2427
+ headers: ta.Optional[ta.Dict[str, str]] = None,
2428
+ ) -> ShellCmd:
2429
+ if content_type is None and json_content:
2430
+ content_type = 'application/json'
2325
2431
 
2326
- #
2432
+ env = {}
2327
2433
 
2328
- no_dependency_cleanup: bool = False
2434
+ header_auth_token: ta.Optional[str]
2435
+ if self._auth_token:
2436
+ header_env_key = self.header_auth_token_env_key
2437
+ env[header_env_key] = self._auth_token
2438
+ header_auth_token = f'${header_env_key}'
2439
+ else:
2440
+ header_auth_token = None
2329
2441
 
2330
- #
2442
+ built_hdrs = self.build_headers(
2443
+ headers,
2444
+ auth_token=header_auth_token,
2445
+ content_type=content_type,
2446
+ )
2331
2447
 
2332
- def __post_init__(self) -> None:
2333
- check.not_isinstance(self.run_options, str)
2448
+ url = f'{self._service_url}/{url}'
2334
2449
 
2335
- def __init__(self, cfg: Config) -> None:
2336
- super().__init__()
2450
+ cmd = ' '.join([
2451
+ 'curl',
2452
+ '-s',
2453
+ '-X', method,
2454
+ url,
2455
+ *[f'-H "{k}: {v}"' for k, v in built_hdrs.items()],
2456
+ ])
2337
2457
 
2338
- self._cfg = cfg
2458
+ return ShellCmd(
2459
+ cmd,
2460
+ env=env,
2461
+ )
2339
2462
 
2340
- self._subprocess_kwargs = {
2341
- **(dict(cwd=self._cfg.cwd) if self._cfg.cwd is not None else {}),
2342
- }
2463
+ def build_post_json_cmd(
2464
+ self,
2465
+ url: str,
2466
+ obj: ta.Any,
2467
+ **kwargs: ta.Any,
2468
+ ) -> ShellCmd:
2469
+ curl_cmd = self.build_cmd(
2470
+ 'POST',
2471
+ url,
2472
+ json_content=True,
2473
+ **kwargs,
2474
+ )
2343
2475
 
2344
- #
2476
+ obj_json = json_dumps_compact(obj)
2345
2477
 
2346
- @property
2347
- def image_tag(self) -> str:
2348
- pfx = 'sha256:'
2349
- if (image := self._cfg.image).startswith(pfx):
2350
- image = image[len(pfx):]
2478
+ return dc.replace(curl_cmd, s=f'{curl_cmd.s} -d {shlex.quote(obj_json)}')
2351
2479
 
2352
- return f'{self._cfg.service}:{image}'
2480
+ #
2353
2481
 
2354
- @cached_nullary
2355
- def tag_image(self) -> str:
2356
- image_tag = self.image_tag
2482
+ @dc.dataclass()
2483
+ class Error(RuntimeError):
2484
+ status_code: int
2485
+ body: ta.Optional[bytes]
2357
2486
 
2358
- subprocesses.check_call(
2359
- 'docker',
2360
- 'tag',
2361
- self._cfg.image,
2362
- image_tag,
2363
- **self._subprocess_kwargs,
2364
- )
2487
+ def __str__(self) -> str:
2488
+ return repr(self)
2365
2489
 
2366
- def delete_tag() -> None:
2367
- subprocesses.check_call(
2368
- 'docker',
2369
- 'rmi',
2370
- image_tag,
2371
- **self._subprocess_kwargs,
2372
- )
2490
+ @dc.dataclass(frozen=True)
2491
+ class Result:
2492
+ status_code: int
2493
+ body: ta.Optional[bytes]
2373
2494
 
2374
- self._enter_context(defer(delete_tag)) # noqa
2495
+ def as_error(self) -> 'GithubServiceCurlClient.Error':
2496
+ return GithubServiceCurlClient.Error(
2497
+ status_code=self.status_code,
2498
+ body=self.body,
2499
+ )
2375
2500
 
2376
- return image_tag
2501
+ def run_cmd(
2502
+ self,
2503
+ cmd: ShellCmd,
2504
+ *,
2505
+ raise_: bool = False,
2506
+ **subprocess_kwargs: ta.Any,
2507
+ ) -> Result:
2508
+ out_file = make_temp_file()
2509
+ with defer(lambda: os.unlink(out_file)):
2510
+ run_cmd = dc.replace(cmd, s=f"{cmd.s} -o {out_file} -w '%{{json}}'")
2377
2511
 
2378
- #
2512
+ out_json_bytes = run_cmd.run(
2513
+ subprocesses.check_output,
2514
+ **subprocess_kwargs,
2515
+ )
2379
2516
 
2380
- def _rewrite_compose_dct(self, in_dct: ta.Dict[str, ta.Any]) -> ta.Dict[str, ta.Any]:
2381
- out = dict(in_dct)
2517
+ out_json = json.loads(out_json_bytes.decode())
2518
+ status_code = check.isinstance(out_json['response_code'], int)
2382
2519
 
2383
- #
2520
+ with open(out_file, 'rb') as f:
2521
+ body = f.read()
2384
2522
 
2385
- in_services = in_dct['services']
2386
- out['services'] = out_services = {}
2523
+ result = self.Result(
2524
+ status_code=status_code,
2525
+ body=body,
2526
+ )
2387
2527
 
2388
- #
2528
+ if raise_ and (500 <= status_code <= 600):
2529
+ raise result.as_error()
2389
2530
 
2390
- in_service: dict = in_services[self._cfg.service]
2391
- out_services[self._cfg.service] = out_service = dict(in_service)
2531
+ return result
2392
2532
 
2393
- out_service['image'] = self.image_tag
2533
+ def run_json_cmd(
2534
+ self,
2535
+ cmd: ShellCmd,
2536
+ *,
2537
+ success_status_codes: ta.Optional[ta.Container[int]] = None,
2538
+ ) -> ta.Optional[ta.Any]:
2539
+ result = self.run_cmd(cmd, raise_=True)
2394
2540
 
2395
- for k in ['build', 'platform']:
2396
- if k in out_service:
2397
- del out_service[k]
2541
+ if success_status_codes is not None:
2542
+ is_success = result.status_code in success_status_codes
2543
+ else:
2544
+ is_success = 200 <= result.status_code < 300
2398
2545
 
2399
- out_service['links'] = [
2400
- f'{l}:{l}' if ':' not in l else l
2401
- for l in out_service.get('links', [])
2402
- ]
2546
+ if is_success:
2547
+ if not (body := result.body):
2548
+ return None
2549
+ return json.loads(body.decode('utf-8-sig'))
2403
2550
 
2404
- #
2551
+ elif result.status_code == 404:
2552
+ return None
2405
2553
 
2406
- depends_on = in_service.get('depends_on', [])
2554
+ else:
2555
+ raise result.as_error()
2407
2556
 
2408
- for dep_service, in_dep_service_dct in list(in_services.items()):
2409
- if dep_service not in depends_on:
2410
- continue
2411
2557
 
2412
- out_dep_service: dict = dict(in_dep_service_dct)
2413
- out_services[dep_service] = out_dep_service
2558
+ ########################################
2559
+ # ../requirements.py
2560
+ """
2561
+ TODO:
2562
+ - pip compile lol
2563
+ - but still support git+ stuff
2564
+ - req.txt format aware hash
2565
+ - more than just whitespace
2566
+ - pyproject req rewriting
2567
+ - download_requirements bootstrap off prev? not worth the dl?
2568
+ - big deps (torch) change less, probably worth it
2569
+ - follow embedded -r automatically like pyp
2570
+ """
2414
2571
 
2415
- out_dep_service['ports'] = []
2416
2572
 
2417
- #
2573
+ ##
2418
2574
 
2419
- return out
2420
2575
 
2421
- @cached_nullary
2422
- def rewrite_compose_file(self) -> str:
2423
- in_dct = read_yaml_file(self._cfg.compose_file)
2576
+ def build_requirements_hash(
2577
+ requirements_txts: ta.Sequence[str],
2578
+ ) -> str:
2579
+ txt_file_contents: dict = {}
2424
2580
 
2425
- out_dct = self._rewrite_compose_dct(in_dct)
2581
+ for txt_file in requirements_txts:
2582
+ txt_file_name = os.path.basename(txt_file)
2583
+ check.not_in(txt_file_name, txt_file_contents)
2584
+ with open(txt_file) as f:
2585
+ txt_contents = f.read()
2586
+ txt_file_contents[txt_file_name] = txt_contents
2426
2587
 
2427
- #
2588
+ #
2589
+
2590
+ lines = []
2591
+ for txt_file, txt_contents in sorted(txt_file_contents.items()):
2592
+ txt_hash = sha256_str(txt_contents)
2593
+ lines.append(f'{txt_file}={txt_hash}')
2594
+
2595
+ return sha256_str('\n'.join(lines))
2596
+
2597
+
2598
+ ##
2599
+
2600
+
2601
+ def download_requirements(
2602
+ image: str,
2603
+ requirements_dir: str,
2604
+ requirements_txts: ta.Sequence[str],
2605
+ ) -> None:
2606
+ requirements_txt_dir = tempfile.mkdtemp()
2607
+ with defer(lambda: shutil.rmtree(requirements_txt_dir)):
2608
+ for rt in requirements_txts:
2609
+ shutil.copyfile(rt, os.path.join(requirements_txt_dir, os.path.basename(rt)))
2610
+
2611
+ subprocesses.check_call(
2612
+ 'docker',
2613
+ 'run',
2614
+ '--rm',
2615
+ '-i',
2616
+ '-v', f'{os.path.abspath(requirements_dir)}:/requirements',
2617
+ '-v', f'{requirements_txt_dir}:/requirements_txt',
2618
+ image,
2619
+ 'pip',
2620
+ 'download',
2621
+ '-d', '/requirements',
2622
+ *itertools.chain.from_iterable(
2623
+ ['-r', f'/requirements_txt/{os.path.basename(rt)}']
2624
+ for rt in requirements_txts
2625
+ ),
2626
+ )
2627
+
2628
+
2629
+ ########################################
2630
+ # ../../../omlish/asyncs/asyncio/subprocesses.py
2631
+
2632
+
2633
+ ##
2634
+
2635
+
2636
+ class AsyncioProcessCommunicator:
2637
+ def __init__(
2638
+ self,
2639
+ proc: asyncio.subprocess.Process,
2640
+ loop: ta.Optional[ta.Any] = None,
2641
+ *,
2642
+ log: ta.Optional[logging.Logger] = None,
2643
+ ) -> None:
2644
+ super().__init__()
2645
+
2646
+ if loop is None:
2647
+ loop = asyncio.get_running_loop()
2648
+
2649
+ self._proc = proc
2650
+ self._loop = loop
2651
+ self._log = log
2652
+
2653
+ self._transport: asyncio.base_subprocess.BaseSubprocessTransport = check.isinstance(
2654
+ proc._transport, # type: ignore # noqa
2655
+ asyncio.base_subprocess.BaseSubprocessTransport,
2656
+ )
2657
+
2658
+ @property
2659
+ def _debug(self) -> bool:
2660
+ return self._loop.get_debug()
2661
+
2662
+ async def _feed_stdin(self, input: bytes) -> None: # noqa
2663
+ stdin = check.not_none(self._proc.stdin)
2664
+ try:
2665
+ if input is not None:
2666
+ stdin.write(input)
2667
+ if self._debug and self._log is not None:
2668
+ self._log.debug('%r communicate: feed stdin (%s bytes)', self, len(input))
2669
+
2670
+ await stdin.drain()
2671
+
2672
+ except (BrokenPipeError, ConnectionResetError) as exc:
2673
+ # communicate() ignores BrokenPipeError and ConnectionResetError. write() and drain() can raise these
2674
+ # exceptions.
2675
+ if self._debug and self._log is not None:
2676
+ self._log.debug('%r communicate: stdin got %r', self, exc)
2677
+
2678
+ if self._debug and self._log is not None:
2679
+ self._log.debug('%r communicate: close stdin', self)
2680
+
2681
+ stdin.close()
2682
+
2683
+ async def _noop(self) -> None:
2684
+ return None
2685
+
2686
+ async def _read_stream(self, fd: int) -> bytes:
2687
+ transport: ta.Any = check.not_none(self._transport.get_pipe_transport(fd))
2688
+
2689
+ if fd == 2:
2690
+ stream = check.not_none(self._proc.stderr)
2691
+ else:
2692
+ check.equal(fd, 1)
2693
+ stream = check.not_none(self._proc.stdout)
2694
+
2695
+ if self._debug and self._log is not None:
2696
+ name = 'stdout' if fd == 1 else 'stderr'
2697
+ self._log.debug('%r communicate: read %s', self, name)
2698
+
2699
+ output = await stream.read()
2700
+
2701
+ if self._debug and self._log is not None:
2702
+ name = 'stdout' if fd == 1 else 'stderr'
2703
+ self._log.debug('%r communicate: close %s', self, name)
2704
+
2705
+ transport.close()
2706
+
2707
+ return output
2708
+
2709
+ class Communication(ta.NamedTuple):
2710
+ stdout: ta.Optional[bytes]
2711
+ stderr: ta.Optional[bytes]
2712
+
2713
+ async def _communicate(
2714
+ self,
2715
+ input: ta.Any = None, # noqa
2716
+ ) -> Communication:
2717
+ stdin_fut: ta.Any
2718
+ if self._proc.stdin is not None:
2719
+ stdin_fut = self._feed_stdin(input)
2720
+ else:
2721
+ stdin_fut = self._noop()
2722
+
2723
+ stdout_fut: ta.Any
2724
+ if self._proc.stdout is not None:
2725
+ stdout_fut = self._read_stream(1)
2726
+ else:
2727
+ stdout_fut = self._noop()
2728
+
2729
+ stderr_fut: ta.Any
2730
+ if self._proc.stderr is not None:
2731
+ stderr_fut = self._read_stream(2)
2732
+ else:
2733
+ stderr_fut = self._noop()
2734
+
2735
+ stdin_res, stdout_res, stderr_res = await asyncio.gather(stdin_fut, stdout_fut, stderr_fut)
2736
+
2737
+ await self._proc.wait()
2738
+
2739
+ return AsyncioProcessCommunicator.Communication(stdout_res, stderr_res)
2740
+
2741
+ async def communicate(
2742
+ self,
2743
+ input: ta.Any = None, # noqa
2744
+ timeout: ta.Optional[float] = None,
2745
+ ) -> Communication:
2746
+ return await asyncio_maybe_timeout(self._communicate(input), timeout)
2747
+
2748
+
2749
+ ##
2750
+
2751
+
2752
+ class AsyncioSubprocesses(AbstractAsyncSubprocesses):
2753
+ async def communicate(
2754
+ self,
2755
+ proc: asyncio.subprocess.Process,
2756
+ input: ta.Any = None, # noqa
2757
+ timeout: ta.Optional[float] = None,
2758
+ ) -> ta.Tuple[ta.Optional[bytes], ta.Optional[bytes]]:
2759
+ return await AsyncioProcessCommunicator(proc).communicate(input, timeout) # noqa
2760
+
2761
+ #
2762
+
2763
+ @contextlib.asynccontextmanager
2764
+ async def popen(
2765
+ self,
2766
+ *cmd: str,
2767
+ shell: bool = False,
2768
+ timeout: ta.Optional[float] = None,
2769
+ **kwargs: ta.Any,
2770
+ ) -> ta.AsyncGenerator[asyncio.subprocess.Process, None]:
2771
+ fac: ta.Any
2772
+ if shell:
2773
+ fac = functools.partial(
2774
+ asyncio.create_subprocess_shell,
2775
+ check.single(cmd),
2776
+ )
2777
+ else:
2778
+ fac = functools.partial(
2779
+ asyncio.create_subprocess_exec,
2780
+ *cmd,
2781
+ )
2782
+
2783
+ with self.prepare_and_wrap( *cmd, shell=shell, **kwargs) as (cmd, kwargs): # noqa
2784
+ proc: asyncio.subprocess.Process = await fac(**kwargs)
2785
+ try:
2786
+ yield proc
2787
+
2788
+ finally:
2789
+ await asyncio_maybe_timeout(proc.wait(), timeout)
2790
+
2791
+ #
2792
+
2793
+ @dc.dataclass(frozen=True)
2794
+ class RunOutput:
2795
+ proc: asyncio.subprocess.Process
2796
+ stdout: ta.Optional[bytes]
2797
+ stderr: ta.Optional[bytes]
2798
+
2799
+ async def run(
2800
+ self,
2801
+ *cmd: str,
2802
+ input: ta.Any = None, # noqa
2803
+ timeout: ta.Optional[float] = None,
2804
+ check: bool = False, # noqa
2805
+ capture_output: ta.Optional[bool] = None,
2806
+ **kwargs: ta.Any,
2807
+ ) -> RunOutput:
2808
+ if capture_output:
2809
+ kwargs.setdefault('stdout', subprocess.PIPE)
2810
+ kwargs.setdefault('stderr', subprocess.PIPE)
2811
+
2812
+ proc: asyncio.subprocess.Process
2813
+ async with self.popen(*cmd, **kwargs) as proc:
2814
+ stdout, stderr = await self.communicate(proc, input, timeout)
2815
+
2816
+ if check and proc.returncode:
2817
+ raise subprocess.CalledProcessError(
2818
+ proc.returncode,
2819
+ cmd,
2820
+ output=stdout,
2821
+ stderr=stderr,
2822
+ )
2823
+
2824
+ return self.RunOutput(
2825
+ proc,
2826
+ stdout,
2827
+ stderr,
2828
+ )
2829
+
2830
+ #
2831
+
2832
+ async def check_call(
2833
+ self,
2834
+ *cmd: str,
2835
+ stdout: ta.Any = sys.stderr,
2836
+ **kwargs: ta.Any,
2837
+ ) -> None:
2838
+ with self.prepare_and_wrap(*cmd, stdout=stdout, check=True, **kwargs) as (cmd, kwargs): # noqa
2839
+ await self.run(*cmd, **kwargs)
2840
+
2841
+ async def check_output(
2842
+ self,
2843
+ *cmd: str,
2844
+ **kwargs: ta.Any,
2845
+ ) -> bytes:
2846
+ with self.prepare_and_wrap(*cmd, stdout=subprocess.PIPE, check=True, **kwargs) as (cmd, kwargs): # noqa
2847
+ return check.not_none((await self.run(*cmd, **kwargs)).stdout)
2848
+
2849
+
2850
+ asyncio_subprocesses = AsyncioSubprocesses()
2851
+
2852
+
2853
+ ########################################
2854
+ # ../compose.py
2855
+ """
2856
+ TODO:
2857
+ - fix rmi - only when not referenced anymore
2858
+ """
2859
+
2860
+
2861
+ ##
2862
+
2863
+
2864
+ def get_compose_service_dependencies(
2865
+ compose_file: str,
2866
+ service: str,
2867
+ ) -> ta.Dict[str, str]:
2868
+ compose_dct = read_yaml_file(compose_file)
2869
+
2870
+ services = compose_dct['services']
2871
+ service_dct = services[service]
2872
+
2873
+ out = {}
2874
+ for dep_service in service_dct.get('depends_on', []):
2875
+ dep_service_dct = services[dep_service]
2876
+ out[dep_service] = dep_service_dct['image']
2877
+
2878
+ return out
2879
+
2880
+
2881
+ ##
2882
+
2883
+
2884
+ class DockerComposeRun(AsyncExitStacked):
2885
+ @dc.dataclass(frozen=True)
2886
+ class Config:
2887
+ compose_file: str
2888
+ service: str
2889
+
2890
+ image: str
2891
+
2892
+ cmd: ShellCmd
2893
+
2894
+ #
2895
+
2896
+ run_options: ta.Optional[ta.Sequence[str]] = None
2897
+
2898
+ cwd: ta.Optional[str] = None
2899
+
2900
+ #
2901
+
2902
+ no_dependencies: bool = False
2903
+ no_dependency_cleanup: bool = False
2904
+
2905
+ #
2906
+
2907
+ def __post_init__(self) -> None:
2908
+ check.not_isinstance(self.run_options, str)
2909
+
2910
+ def __init__(self, cfg: Config) -> None:
2911
+ super().__init__()
2912
+
2913
+ self._cfg = cfg
2914
+
2915
+ self._subprocess_kwargs = {
2916
+ **(dict(cwd=self._cfg.cwd) if self._cfg.cwd is not None else {}),
2917
+ }
2918
+
2919
+ #
2920
+
2921
+ def _rewrite_compose_dct(self, in_dct: ta.Dict[str, ta.Any]) -> ta.Dict[str, ta.Any]:
2922
+ out = dict(in_dct)
2923
+
2924
+ #
2925
+
2926
+ in_services = in_dct['services']
2927
+ out['services'] = out_services = {}
2928
+
2929
+ #
2930
+
2931
+ in_service: dict = in_services[self._cfg.service]
2932
+ out_services[self._cfg.service] = out_service = dict(in_service)
2933
+
2934
+ out_service['image'] = self._cfg.image
2935
+
2936
+ for k in ['build', 'platform']:
2937
+ if k in out_service:
2938
+ del out_service[k]
2939
+
2940
+ out_service['links'] = [
2941
+ f'{l}:{l}' if ':' not in l else l
2942
+ for l in out_service.get('links', [])
2943
+ ]
2944
+
2945
+ #
2946
+
2947
+ if not self._cfg.no_dependencies:
2948
+ depends_on = in_service.get('depends_on', [])
2949
+
2950
+ for dep_service, in_dep_service_dct in list(in_services.items()):
2951
+ if dep_service not in depends_on:
2952
+ continue
2953
+
2954
+ out_dep_service: dict = dict(in_dep_service_dct)
2955
+ out_services[dep_service] = out_dep_service
2956
+
2957
+ out_dep_service['ports'] = []
2958
+
2959
+ else:
2960
+ out_service['depends_on'] = []
2961
+ out_service['links'] = []
2962
+
2963
+ #
2964
+
2965
+ return out
2966
+
2967
+ @cached_nullary
2968
+ def rewrite_compose_file(self) -> str:
2969
+ in_dct = read_yaml_file(self._cfg.compose_file)
2970
+
2971
+ out_dct = self._rewrite_compose_dct(in_dct)
2972
+
2973
+ #
2428
2974
 
2429
2975
  out_compose_file = make_temp_file()
2430
2976
  self._enter_context(defer(lambda: os.unlink(out_compose_file))) # noqa
@@ -2438,22 +2984,20 @@ class DockerComposeRun(ExitStacked):
2438
2984
 
2439
2985
  #
2440
2986
 
2441
- def _cleanup_dependencies(self) -> None:
2442
- subprocesses.check_call(
2987
+ async def _cleanup_dependencies(self) -> None:
2988
+ await asyncio_subprocesses.check_call(
2443
2989
  'docker',
2444
2990
  'compose',
2445
2991
  '-f', self.rewrite_compose_file(),
2446
2992
  'down',
2447
2993
  )
2448
2994
 
2449
- def run(self) -> None:
2450
- self.tag_image()
2451
-
2995
+ async def run(self) -> None:
2452
2996
  compose_file = self.rewrite_compose_file()
2453
2997
 
2454
- with contextlib.ExitStack() as es:
2455
- if not self._cfg.no_dependency_cleanup:
2456
- es.enter_context(defer(self._cleanup_dependencies)) # noqa
2998
+ async with contextlib.AsyncExitStack() as es:
2999
+ if not (self._cfg.no_dependencies or self._cfg.no_dependency_cleanup):
3000
+ await es.enter_async_context(adefer(self._cleanup_dependencies)) # noqa
2457
3001
 
2458
3002
  sh_cmd = ' '.join([
2459
3003
  'docker',
@@ -2472,8 +3016,8 @@ class DockerComposeRun(ExitStacked):
2472
3016
 
2473
3017
  run_cmd = dc.replace(self._cfg.cmd, s=sh_cmd)
2474
3018
 
2475
- run_cmd.run(
2476
- subprocesses.check_call,
3019
+ await run_cmd.run(
3020
+ asyncio_subprocesses.check_call,
2477
3021
  **self._subprocess_kwargs,
2478
3022
  )
2479
3023
 
@@ -2525,8 +3069,8 @@ def read_docker_tar_image_id(tar_file: str) -> str:
2525
3069
  ##
2526
3070
 
2527
3071
 
2528
- def is_docker_image_present(image: str) -> bool:
2529
- out = subprocesses.check_output(
3072
+ async def is_docker_image_present(image: str) -> bool:
3073
+ out = await asyncio_subprocesses.check_output(
2530
3074
  'docker',
2531
3075
  'images',
2532
3076
  '--format', 'json',
@@ -2541,333 +3085,263 @@ def is_docker_image_present(image: str) -> bool:
2541
3085
  return True
2542
3086
 
2543
3087
 
2544
- def pull_docker_image(
3088
+ async def pull_docker_image(
2545
3089
  image: str,
2546
3090
  ) -> None:
2547
- subprocesses.check_call(
3091
+ await asyncio_subprocesses.check_call(
2548
3092
  'docker',
2549
3093
  'pull',
2550
3094
  image,
2551
3095
  )
2552
3096
 
2553
3097
 
2554
- def build_docker_image(
3098
+ async def build_docker_image(
2555
3099
  docker_file: str,
2556
3100
  *,
3101
+ tag: ta.Optional[str] = None,
2557
3102
  cwd: ta.Optional[str] = None,
2558
3103
  ) -> str:
2559
3104
  id_file = make_temp_file()
2560
3105
  with defer(lambda: os.unlink(id_file)):
2561
- subprocesses.check_call(
3106
+ await asyncio_subprocesses.check_call(
2562
3107
  'docker',
2563
3108
  'build',
2564
3109
  '-f', os.path.abspath(docker_file),
2565
3110
  '--iidfile', id_file,
2566
3111
  '--squash',
3112
+ *(['--tag', tag] if tag is not None else []),
2567
3113
  '.',
2568
3114
  **(dict(cwd=cwd) if cwd is not None else {}),
2569
3115
  )
2570
3116
 
2571
- with open(id_file) as f:
3117
+ with open(id_file) as f: # noqa
2572
3118
  image_id = check.single(f.read().strip().splitlines()).strip()
2573
3119
 
2574
3120
  return image_id
2575
3121
 
2576
3122
 
3123
+ async def tag_docker_image(image: str, tag: str) -> None:
3124
+ await asyncio_subprocesses.check_call(
3125
+ 'docker',
3126
+ 'tag',
3127
+ image,
3128
+ tag,
3129
+ )
3130
+
3131
+
3132
+ async def delete_docker_tag(tag: str) -> None:
3133
+ await asyncio_subprocesses.check_call(
3134
+ 'docker',
3135
+ 'rmi',
3136
+ tag,
3137
+ )
3138
+
3139
+
2577
3140
  ##
2578
3141
 
2579
3142
 
2580
- def save_docker_tar_cmd(
3143
+ async def save_docker_tar_cmd(
2581
3144
  image: str,
2582
3145
  output_cmd: ShellCmd,
2583
3146
  ) -> None:
2584
3147
  cmd = dc.replace(output_cmd, s=f'docker save {image} | {output_cmd.s}')
2585
- cmd.run(subprocesses.check_call)
3148
+ await cmd.run(asyncio_subprocesses.check_call)
2586
3149
 
2587
3150
 
2588
- def save_docker_tar(
3151
+ async def save_docker_tar(
2589
3152
  image: str,
2590
3153
  tar_file: str,
2591
3154
  ) -> None:
2592
- return save_docker_tar_cmd(
3155
+ return await save_docker_tar_cmd(
2593
3156
  image,
2594
3157
  ShellCmd(f'cat > {shlex.quote(tar_file)}'),
2595
3158
  )
2596
3159
 
2597
3160
 
2598
- #
2599
-
2600
-
2601
- def load_docker_tar_cmd(
2602
- input_cmd: ShellCmd,
2603
- ) -> str:
2604
- cmd = dc.replace(input_cmd, s=f'{input_cmd.s} | docker load')
2605
-
2606
- out = cmd.run(subprocesses.check_output).decode()
2607
-
2608
- line = check.single(out.strip().splitlines())
2609
- loaded = line.partition(':')[2].strip()
2610
- return loaded
2611
-
2612
-
2613
- def load_docker_tar(
2614
- tar_file: str,
2615
- ) -> str:
2616
- return load_docker_tar_cmd(ShellCmd(f'cat {shlex.quote(tar_file)}'))
2617
-
2618
-
2619
- ########################################
2620
- # ../github/cache.py
2621
-
2622
-
2623
- ##
2624
-
2625
-
2626
- class GithubV1CacheShellClient:
2627
- BASE_URL_ENV_KEY = 'ACTIONS_CACHE_URL'
2628
- AUTH_TOKEN_ENV_KEY = 'ACTIONS_RUNTIME_TOKEN' # noqa
2629
-
2630
- def __init__(
2631
- self,
2632
- *,
2633
- base_url: ta.Optional[str] = None,
2634
- auth_token: ta.Optional[str] = None,
2635
- ) -> None:
2636
- super().__init__()
2637
-
2638
- if base_url is None:
2639
- base_url = os.environ[self.BASE_URL_ENV_KEY]
2640
- self._base_url = check.non_empty_str(base_url)
2641
-
2642
- if auth_token is None:
2643
- auth_token = os.environ.get(self.AUTH_TOKEN_ENV_KEY)
2644
- self._auth_token = auth_token
2645
-
2646
- self._service_url = GithubCacheServiceV1.get_service_url(self._base_url)
2647
-
2648
- #
2649
-
2650
- _MISSING = object()
2651
-
2652
- def build_headers(
2653
- self,
2654
- *,
2655
- auth_token: ta.Any = _MISSING,
2656
- content_type: ta.Optional[str] = None,
2657
- ) -> ta.Dict[str, str]:
2658
- dct = {
2659
- 'Accept': f'application/json;api-version={GithubCacheServiceV1.API_VERSION}',
2660
- }
3161
+ #
2661
3162
 
2662
- if auth_token is self._MISSING:
2663
- auth_token = self._auth_token
2664
- if auth_token:
2665
- dct['Authorization'] = f'Bearer {auth_token}'
2666
3163
 
2667
- if content_type is not None:
2668
- dct['Content-Type'] = content_type
3164
+ async def load_docker_tar_cmd(
3165
+ input_cmd: ShellCmd,
3166
+ ) -> str:
3167
+ cmd = dc.replace(input_cmd, s=f'{input_cmd.s} | docker load')
2669
3168
 
2670
- return dct
3169
+ out = (await cmd.run(asyncio_subprocesses.check_output)).decode()
2671
3170
 
2672
- #
3171
+ line = check.single(out.strip().splitlines())
3172
+ loaded = line.partition(':')[2].strip()
3173
+ return loaded
2673
3174
 
2674
- HEADER_AUTH_TOKEN_ENV_KEY = '_GITHUB_CACHE_AUTH_TOKEN' # noqa
2675
3175
 
2676
- def build_curl_cmd(
2677
- self,
2678
- method: str,
2679
- url: str,
2680
- *,
2681
- json_content: bool = False,
2682
- content_type: ta.Optional[str] = None,
2683
- ) -> ShellCmd:
2684
- if content_type is None and json_content:
2685
- content_type = 'application/json'
3176
+ async def load_docker_tar(
3177
+ tar_file: str,
3178
+ ) -> str:
3179
+ return await load_docker_tar_cmd(ShellCmd(f'cat {shlex.quote(tar_file)}'))
2686
3180
 
2687
- env = {}
2688
3181
 
2689
- header_auth_token: ta.Optional[str]
2690
- if self._auth_token:
2691
- env[self.HEADER_AUTH_TOKEN_ENV_KEY] = self._auth_token
2692
- header_auth_token = f'${self.HEADER_AUTH_TOKEN_ENV_KEY}'
2693
- else:
2694
- header_auth_token = None
3182
+ ########################################
3183
+ # ../github/cache.py
2695
3184
 
2696
- hdrs = self.build_headers(
2697
- auth_token=header_auth_token,
2698
- content_type=content_type,
2699
- )
2700
3185
 
2701
- url = f'{self._service_url}/{url}'
3186
+ ##
2702
3187
 
2703
- cmd = ' '.join([
2704
- 'curl',
2705
- '-s',
2706
- '-X', method,
2707
- url,
2708
- *[f'-H "{k}: {v}"' for k, v in hdrs.items()],
2709
- ])
2710
3188
 
2711
- return ShellCmd(
2712
- cmd,
2713
- env=env,
2714
- )
3189
+ class GithubCacheShellClient(abc.ABC):
3190
+ class Entry(abc.ABC): # noqa
3191
+ pass
2715
3192
 
2716
- def build_post_json_curl_cmd(
2717
- self,
2718
- url: str,
2719
- obj: ta.Any,
2720
- **kwargs: ta.Any,
2721
- ) -> ShellCmd:
2722
- curl_cmd = self.build_curl_cmd(
2723
- 'POST',
2724
- url,
2725
- json_content=True,
2726
- **kwargs,
2727
- )
3193
+ @abc.abstractmethod
3194
+ def run_get_entry(self, key: str) -> ta.Optional[Entry]:
3195
+ raise NotImplementedError
2728
3196
 
2729
- obj_json = json_dumps_compact(obj)
3197
+ @abc.abstractmethod
3198
+ def download_get_entry(self, entry: Entry, out_file: str) -> None:
3199
+ raise NotImplementedError
2730
3200
 
2731
- return dc.replace(curl_cmd, s=f'{curl_cmd.s} -d {shlex.quote(obj_json)}')
3201
+ @abc.abstractmethod
3202
+ def upload_cache_entry(self, key: str, in_file: str) -> None:
3203
+ raise NotImplementedError
2732
3204
 
2733
- #
2734
3205
 
2735
- @dc.dataclass()
2736
- class CurlError(RuntimeError):
2737
- status_code: int
2738
- body: ta.Optional[bytes]
3206
+ #
2739
3207
 
2740
- def __str__(self) -> str:
2741
- return repr(self)
2742
3208
 
2743
- @dc.dataclass(frozen=True)
2744
- class CurlResult:
2745
- status_code: int
2746
- body: ta.Optional[bytes]
3209
+ class GithubCacheServiceV1ShellClient(GithubCacheShellClient):
3210
+ BASE_URL_ENV_KEY = 'ACTIONS_CACHE_URL'
3211
+ AUTH_TOKEN_ENV_KEY = 'ACTIONS_RUNTIME_TOKEN' # noqa
2747
3212
 
2748
- def as_error(self) -> 'GithubV1CacheShellClient.CurlError':
2749
- return GithubV1CacheShellClient.CurlError(
2750
- status_code=self.status_code,
2751
- body=self.body,
2752
- )
3213
+ KEY_SUFFIX_ENV_KEY = 'GITHUB_RUN_ID'
3214
+
3215
+ CACHE_VERSION: ta.ClassVar[int] = 1
3216
+
3217
+ #
2753
3218
 
2754
- def run_curl_cmd(
3219
+ def __init__(
2755
3220
  self,
2756
- cmd: ShellCmd,
2757
3221
  *,
2758
- raise_: bool = False,
2759
- ) -> CurlResult:
2760
- out_file = make_temp_file()
2761
- with defer(lambda: os.unlink(out_file)):
2762
- run_cmd = dc.replace(cmd, s=f"{cmd.s} -o {out_file} -w '%{{json}}'")
3222
+ base_url: ta.Optional[str] = None,
3223
+ auth_token: ta.Optional[str] = None,
2763
3224
 
2764
- out_json_bytes = run_cmd.run(subprocesses.check_output)
3225
+ key_prefix: ta.Optional[str] = None,
3226
+ key_suffix: ta.Optional[str] = None,
3227
+ ) -> None:
3228
+ super().__init__()
2765
3229
 
2766
- out_json = json.loads(out_json_bytes.decode())
2767
- status_code = check.isinstance(out_json['response_code'], int)
3230
+ #
2768
3231
 
2769
- with open(out_file, 'rb') as f:
2770
- body = f.read()
3232
+ if base_url is None:
3233
+ base_url = os.environ[self.BASE_URL_ENV_KEY]
3234
+ service_url = GithubCacheServiceV1.get_service_url(base_url)
2771
3235
 
2772
- result = self.CurlResult(
2773
- status_code=status_code,
2774
- body=body,
2775
- )
3236
+ if auth_token is None:
3237
+ auth_token = os.environ.get(self.AUTH_TOKEN_ENV_KEY)
2776
3238
 
2777
- if raise_ and (500 <= status_code <= 600):
2778
- raise result.as_error()
3239
+ self._curl = GithubServiceCurlClient(
3240
+ service_url,
3241
+ auth_token,
3242
+ api_version=GithubCacheServiceV1.API_VERSION,
3243
+ )
2779
3244
 
2780
- return result
3245
+ #
2781
3246
 
2782
- def run_json_curl_cmd(
2783
- self,
2784
- cmd: ShellCmd,
2785
- *,
2786
- success_status_codes: ta.Optional[ta.Container[int]] = None,
2787
- ) -> ta.Optional[ta.Any]:
2788
- result = self.run_curl_cmd(cmd, raise_=True)
3247
+ self._key_prefix = key_prefix
2789
3248
 
2790
- if success_status_codes is not None:
2791
- is_success = result.status_code in success_status_codes
2792
- else:
2793
- is_success = 200 <= result.status_code < 300
3249
+ if key_suffix is None:
3250
+ key_suffix = os.environ[self.KEY_SUFFIX_ENV_KEY]
3251
+ self._key_suffix = check.non_empty_str(key_suffix)
2794
3252
 
2795
- if is_success:
2796
- if not (body := result.body):
2797
- return None
2798
- return json.loads(body.decode('utf-8-sig'))
3253
+ #
2799
3254
 
2800
- elif result.status_code == 404:
2801
- return None
3255
+ KEY_PART_SEPARATOR = '--'
2802
3256
 
2803
- else:
2804
- raise result.as_error()
3257
+ def fix_key(self, s: str) -> str:
3258
+ return self.KEY_PART_SEPARATOR.join([
3259
+ *([self._key_prefix] if self._key_prefix else []),
3260
+ s,
3261
+ self._key_suffix,
3262
+ ])
3263
+
3264
+ #
3265
+
3266
+ @dc.dataclass(frozen=True)
3267
+ class Entry(GithubCacheShellClient.Entry):
3268
+ artifact: GithubCacheServiceV1.ArtifactCacheEntry
2805
3269
 
2806
3270
  #
2807
3271
 
2808
3272
  def build_get_entry_curl_cmd(self, key: str) -> ShellCmd:
2809
- return self.build_curl_cmd(
3273
+ fixed_key = self.fix_key(key)
3274
+
3275
+ qp = dict(
3276
+ keys=fixed_key,
3277
+ version=str(self.CACHE_VERSION),
3278
+ )
3279
+
3280
+ return self._curl.build_cmd(
2810
3281
  'GET',
2811
- f'cache?keys={key}',
3282
+ shlex.quote('?'.join([
3283
+ 'cache',
3284
+ '&'.join([
3285
+ f'{k}={urllib.parse.quote_plus(v)}'
3286
+ for k, v in qp.items()
3287
+ ]),
3288
+ ])),
2812
3289
  )
2813
3290
 
2814
- def run_get_entry(self, key: str) -> ta.Optional[GithubCacheServiceV1.ArtifactCacheEntry]:
2815
- curl_cmd = self.build_get_entry_curl_cmd(key)
3291
+ def run_get_entry(self, key: str) -> ta.Optional[Entry]:
3292
+ fixed_key = self.fix_key(key)
3293
+ curl_cmd = self.build_get_entry_curl_cmd(fixed_key)
2816
3294
 
2817
- obj = self.run_json_curl_cmd(
3295
+ obj = self._curl.run_json_cmd(
2818
3296
  curl_cmd,
2819
3297
  success_status_codes=[200, 204],
2820
3298
  )
2821
3299
  if obj is None:
2822
3300
  return None
2823
3301
 
2824
- return GithubCacheServiceV1.dataclass_from_json(
3302
+ return self.Entry(GithubCacheServiceV1.dataclass_from_json(
2825
3303
  GithubCacheServiceV1.ArtifactCacheEntry,
2826
3304
  obj,
2827
- )
3305
+ ))
2828
3306
 
2829
3307
  #
2830
3308
 
2831
- def build_download_get_entry_cmd(
2832
- self,
2833
- entry: GithubCacheServiceV1.ArtifactCacheEntry,
2834
- out_file: str,
2835
- ) -> ShellCmd:
3309
+ def build_download_get_entry_cmd(self, entry: Entry, out_file: str) -> ShellCmd:
2836
3310
  return ShellCmd(' '.join([
2837
3311
  'aria2c',
2838
3312
  '-x', '4',
2839
3313
  '-o', out_file,
2840
- check.non_empty_str(entry.archive_location),
3314
+ check.non_empty_str(entry.artifact.archive_location),
2841
3315
  ]))
2842
3316
 
2843
- def download_get_entry(
2844
- self,
2845
- entry: GithubCacheServiceV1.ArtifactCacheEntry,
2846
- out_file: str,
2847
- ) -> None:
2848
- dl_cmd = self.build_download_get_entry_cmd(entry, out_file)
3317
+ def download_get_entry(self, entry: GithubCacheShellClient.Entry, out_file: str) -> None:
3318
+ dl_cmd = self.build_download_get_entry_cmd(
3319
+ check.isinstance(entry, GithubCacheServiceV1ShellClient.Entry),
3320
+ out_file,
3321
+ )
2849
3322
  dl_cmd.run(subprocesses.check_call)
2850
3323
 
2851
3324
  #
2852
3325
 
2853
- def upload_cache_entry(
2854
- self,
2855
- key: str,
2856
- in_file: str,
2857
- ) -> None:
3326
+ def upload_cache_entry(self, key: str, in_file: str) -> None:
3327
+ fixed_key = self.fix_key(key)
3328
+
2858
3329
  check.state(os.path.isfile(in_file))
2859
3330
 
2860
3331
  file_size = os.stat(in_file).st_size
2861
3332
 
3333
+ #
3334
+
2862
3335
  reserve_req = GithubCacheServiceV1.ReserveCacheRequest(
2863
- key=key,
3336
+ key=fixed_key,
2864
3337
  cache_size=file_size,
3338
+ version=str(self.CACHE_VERSION),
2865
3339
  )
2866
- reserve_cmd = self.build_post_json_curl_cmd(
3340
+ reserve_cmd = self._curl.build_post_json_cmd(
2867
3341
  'caches',
2868
3342
  GithubCacheServiceV1.dataclass_to_json(reserve_req),
2869
3343
  )
2870
- reserve_resp_obj: ta.Any = check.not_none(self.run_json_curl_cmd(
3344
+ reserve_resp_obj: ta.Any = check.not_none(self._curl.run_json_cmd(
2871
3345
  reserve_cmd,
2872
3346
  success_status_codes=[201],
2873
3347
  ))
@@ -2875,8 +3349,66 @@ class GithubV1CacheShellClient:
2875
3349
  GithubCacheServiceV1.ReserveCacheResponse,
2876
3350
  reserve_resp_obj,
2877
3351
  )
3352
+ cache_id = check.isinstance(reserve_resp.cache_id, int)
2878
3353
 
2879
- raise NotImplementedError
3354
+ #
3355
+
3356
+ tmp_file = make_temp_file()
3357
+
3358
+ print(f'{file_size=}')
3359
+ num_written = 0
3360
+ chunk_size = 32 * 1024 * 1024
3361
+ for i in range((file_size // chunk_size) + (1 if file_size % chunk_size else 0)):
3362
+ ofs = i * chunk_size
3363
+ sz = min(chunk_size, file_size - ofs)
3364
+
3365
+ patch_cmd = self._curl.build_cmd(
3366
+ 'PATCH',
3367
+ f'caches/{cache_id}',
3368
+ content_type='application/octet-stream',
3369
+ headers={
3370
+ 'Content-Range': f'bytes {ofs}-{ofs + sz - 1}/*',
3371
+ },
3372
+ )
3373
+
3374
+ #
3375
+
3376
+ # patch_data_cmd = dc.replace(patch_cmd, s=' | '.join([
3377
+ # f'dd if={in_file} bs={chunk_size} skip={i} count=1 status=none',
3378
+ # f'{patch_cmd.s} --data-binary -',
3379
+ # ]))
3380
+ # print(f'{patch_data_cmd.s=}')
3381
+ # patch_result = self._curl.run_cmd(patch_data_cmd, raise_=True)
3382
+
3383
+ #
3384
+
3385
+ with open(in_file, 'rb') as f:
3386
+ f.seek(ofs)
3387
+ buf = f.read(sz)
3388
+ with open(tmp_file, 'wb') as f:
3389
+ f.write(buf)
3390
+ num_written += len(buf)
3391
+ print(f'{num_written=}')
3392
+ patch_data_cmd = dc.replace(patch_cmd, s=f'{patch_cmd.s} --data-binary @{tmp_file}')
3393
+ print(f'{patch_data_cmd.s=}')
3394
+ patch_result = self._curl.run_cmd(patch_data_cmd, raise_=True)
3395
+
3396
+ #
3397
+
3398
+ check.equal(patch_result.status_code, 204)
3399
+ ofs += sz
3400
+
3401
+ #
3402
+
3403
+ commit_req = GithubCacheServiceV1.CommitCacheRequest(
3404
+ size=file_size,
3405
+ )
3406
+ commit_cmd = self._curl.build_post_json_cmd(
3407
+ f'caches/{cache_id}',
3408
+ GithubCacheServiceV1.dataclass_to_json(commit_req),
3409
+ )
3410
+ commit_result = self._curl.run_cmd(commit_cmd, raise_=True)
3411
+ check.equal(commit_result.status_code, 204)
2880
3412
 
2881
3413
 
2882
3414
  ##
@@ -2887,15 +3419,15 @@ class GithubShellCache(ShellCache):
2887
3419
  self,
2888
3420
  dir: str, # noqa
2889
3421
  *,
2890
- client: ta.Optional[GithubV1CacheShellClient] = None,
3422
+ client: ta.Optional[GithubCacheShellClient] = None,
2891
3423
  ) -> None:
2892
3424
  super().__init__()
2893
3425
 
2894
3426
  self._dir = check.not_none(dir)
2895
3427
 
2896
3428
  if client is None:
2897
- client = GithubV1CacheShellClient()
2898
- self._client = client
3429
+ client = GithubCacheServiceV1ShellClient()
3430
+ self._client: GithubCacheShellClient = client
2899
3431
 
2900
3432
  self._local = DirectoryFileCache(self._dir)
2901
3433
 
@@ -2957,82 +3489,11 @@ class GithubShellCache(ShellCache):
2957
3489
  )
2958
3490
 
2959
3491
 
2960
- ########################################
2961
- # ../requirements.py
2962
- """
2963
- TODO:
2964
- - pip compile lol
2965
- - but still support git+ stuff
2966
- - req.txt format aware hash
2967
- - more than just whitespace
2968
- - pyproject req rewriting
2969
- - download_requirements bootstrap off prev? not worth the dl?
2970
- - big deps (torch) change less, probably worth it
2971
- - follow embedded -r automatically like pyp
2972
- """
2973
-
2974
-
2975
- ##
2976
-
2977
-
2978
- def build_requirements_hash(
2979
- requirements_txts: ta.Sequence[str],
2980
- ) -> str:
2981
- txt_file_contents: dict = {}
2982
-
2983
- for txt_file in requirements_txts:
2984
- txt_file_name = os.path.basename(txt_file)
2985
- check.not_in(txt_file_name, txt_file_contents)
2986
- with open(txt_file) as f:
2987
- txt_contents = f.read()
2988
- txt_file_contents[txt_file_name] = txt_contents
2989
-
2990
- #
2991
-
2992
- lines = []
2993
- for txt_file, txt_contents in sorted(txt_file_contents.items()):
2994
- txt_hash = sha256_str(txt_contents)
2995
- lines.append(f'{txt_file}={txt_hash}')
2996
-
2997
- return sha256_str('\n'.join(lines))
2998
-
2999
-
3000
- ##
3001
-
3002
-
3003
- def download_requirements(
3004
- image: str,
3005
- requirements_dir: str,
3006
- requirements_txts: ta.Sequence[str],
3007
- ) -> None:
3008
- requirements_txt_dir = tempfile.mkdtemp()
3009
- with defer(lambda: shutil.rmtree(requirements_txt_dir)):
3010
- for rt in requirements_txts:
3011
- shutil.copyfile(rt, os.path.join(requirements_txt_dir, os.path.basename(rt)))
3012
-
3013
- subprocesses.check_call(
3014
- 'docker',
3015
- 'run',
3016
- '--rm',
3017
- '-i',
3018
- '-v', f'{os.path.abspath(requirements_dir)}:/requirements',
3019
- '-v', f'{requirements_txt_dir}:/requirements_txt',
3020
- image,
3021
- 'pip',
3022
- 'download',
3023
- '-d', '/requirements',
3024
- *itertools.chain.from_iterable(
3025
- ['-r', f'/requirements_txt/{os.path.basename(rt)}']
3026
- for rt in requirements_txts
3027
- ),
3028
- )
3029
-
3030
-
3031
3492
  ########################################
3032
3493
  # ../ci.py
3033
3494
 
3034
3495
 
3035
- class Ci(ExitStacked):
3496
+ class Ci(AsyncExitStacked):
3036
3497
  FILE_NAME_HASH_LEN = 16
3037
3498
 
3038
3499
  @dc.dataclass(frozen=True)
@@ -3049,6 +3510,9 @@ class Ci(ExitStacked):
3049
3510
  requirements_txts: ta.Optional[ta.Sequence[str]] = None
3050
3511
 
3051
3512
  always_pull: bool = False
3513
+ always_build: bool = False
3514
+
3515
+ no_dependencies: bool = False
3052
3516
 
3053
3517
  def __post_init__(self) -> None:
3054
3518
  check.not_isinstance(self.requirements_txts, str)
@@ -3068,7 +3532,7 @@ class Ci(ExitStacked):
3068
3532
 
3069
3533
  #
3070
3534
 
3071
- def _load_cache_docker_image(self, key: str) -> ta.Optional[str]:
3535
+ async def _load_cache_docker_image(self, key: str) -> ta.Optional[str]:
3072
3536
  if self._shell_cache is None:
3073
3537
  return None
3074
3538
 
@@ -3078,9 +3542,9 @@ class Ci(ExitStacked):
3078
3542
 
3079
3543
  get_cache_cmd = dc.replace(get_cache_cmd, s=f'{get_cache_cmd.s} | zstd -cd --long') # noqa
3080
3544
 
3081
- return load_docker_tar_cmd(get_cache_cmd)
3545
+ return await load_docker_tar_cmd(get_cache_cmd)
3082
3546
 
3083
- def _save_cache_docker_image(self, key: str, image: str) -> None:
3547
+ async def _save_cache_docker_image(self, key: str, image: str) -> None:
3084
3548
  if self._shell_cache is None:
3085
3549
  return
3086
3550
 
@@ -3089,12 +3553,12 @@ class Ci(ExitStacked):
3089
3553
 
3090
3554
  put_cache_cmd = dc.replace(put_cache_cmd, s=f'zstd | {put_cache_cmd.s}')
3091
3555
 
3092
- save_docker_tar_cmd(image, put_cache_cmd)
3556
+ await save_docker_tar_cmd(image, put_cache_cmd)
3093
3557
 
3094
3558
  #
3095
3559
 
3096
- def _load_docker_image(self, image: str) -> None:
3097
- if not self._cfg.always_pull and is_docker_image_present(image):
3560
+ async def _load_docker_image(self, image: str) -> None:
3561
+ if not self._cfg.always_pull and (await is_docker_image_present(image)):
3098
3562
  return
3099
3563
 
3100
3564
  dep_suffix = image
@@ -3102,63 +3566,79 @@ class Ci(ExitStacked):
3102
3566
  dep_suffix = dep_suffix.replace(c, '-')
3103
3567
 
3104
3568
  cache_key = f'docker-{dep_suffix}'
3105
- if self._load_cache_docker_image(cache_key) is not None:
3569
+ if (await self._load_cache_docker_image(cache_key)) is not None:
3106
3570
  return
3107
3571
 
3108
- pull_docker_image(image)
3572
+ await pull_docker_image(image)
3109
3573
 
3110
- self._save_cache_docker_image(cache_key, image)
3574
+ await self._save_cache_docker_image(cache_key, image)
3111
3575
 
3112
- def load_docker_image(self, image: str) -> None:
3576
+ async def load_docker_image(self, image: str) -> None:
3113
3577
  with log_timing_context(f'Load docker image: {image}'):
3114
- self._load_docker_image(image)
3578
+ await self._load_docker_image(image)
3115
3579
 
3116
- @cached_nullary
3117
- def load_compose_service_dependencies(self) -> None:
3580
+ @async_cached_nullary
3581
+ async def load_compose_service_dependencies(self) -> None:
3118
3582
  deps = get_compose_service_dependencies(
3119
3583
  self._cfg.compose_file,
3120
3584
  self._cfg.service,
3121
3585
  )
3122
3586
 
3123
3587
  for dep_image in deps.values():
3124
- self.load_docker_image(dep_image)
3588
+ await self.load_docker_image(dep_image)
3125
3589
 
3126
3590
  #
3127
3591
 
3128
- def _resolve_ci_image(self) -> str:
3129
- docker_file_hash = build_docker_file_hash(self._cfg.docker_file)[:self.FILE_NAME_HASH_LEN]
3592
+ @cached_nullary
3593
+ def docker_file_hash(self) -> str:
3594
+ return build_docker_file_hash(self._cfg.docker_file)[:self.FILE_NAME_HASH_LEN]
3595
+
3596
+ async def _resolve_ci_image(self) -> str:
3597
+ cache_key = f'ci-{self.docker_file_hash()}'
3598
+ image_tag = f'{self._cfg.service}:{cache_key}'
3599
+
3600
+ if not self._cfg.always_build and (await is_docker_image_present(image_tag)):
3601
+ return image_tag
3130
3602
 
3131
- cache_key = f'ci-{docker_file_hash}'
3132
- if (cache_image_id := self._load_cache_docker_image(cache_key)) is not None:
3133
- return cache_image_id
3603
+ if (cache_image_id := await self._load_cache_docker_image(cache_key)) is not None:
3604
+ await tag_docker_image(
3605
+ cache_image_id,
3606
+ image_tag,
3607
+ )
3608
+ return image_tag
3134
3609
 
3135
- image_id = build_docker_image(
3610
+ image_id = await build_docker_image(
3136
3611
  self._cfg.docker_file,
3612
+ tag=image_tag,
3137
3613
  cwd=self._cfg.project_dir,
3138
3614
  )
3139
3615
 
3140
- self._save_cache_docker_image(cache_key, image_id)
3616
+ await self._save_cache_docker_image(cache_key, image_id)
3141
3617
 
3142
- return image_id
3618
+ return image_tag
3143
3619
 
3144
- @cached_nullary
3145
- def resolve_ci_image(self) -> str:
3620
+ @async_cached_nullary
3621
+ async def resolve_ci_image(self) -> str:
3146
3622
  with log_timing_context('Resolve ci image') as ltc:
3147
- image_id = self._resolve_ci_image()
3623
+ image_id = await self._resolve_ci_image()
3148
3624
  ltc.set_description(f'Resolve ci image: {image_id}')
3149
3625
  return image_id
3150
3626
 
3151
3627
  #
3152
3628
 
3153
- def _resolve_requirements_dir(self) -> str:
3154
- requirements_txts = [
3629
+ @cached_nullary
3630
+ def requirements_txts(self) -> ta.Sequence[str]:
3631
+ return [
3155
3632
  os.path.join(self._cfg.project_dir, rf)
3156
3633
  for rf in check.not_none(self._cfg.requirements_txts)
3157
3634
  ]
3158
3635
 
3159
- requirements_hash = build_requirements_hash(requirements_txts)[:self.FILE_NAME_HASH_LEN]
3636
+ @cached_nullary
3637
+ def requirements_hash(self) -> str:
3638
+ return build_requirements_hash(self.requirements_txts())[:self.FILE_NAME_HASH_LEN]
3160
3639
 
3161
- tar_file_key = f'requirements-{requirements_hash}'
3640
+ async def _resolve_requirements_dir(self) -> str:
3641
+ tar_file_key = f'requirements-{self.docker_file_hash()}-{self.requirements_hash()}'
3162
3642
  tar_file_name = f'{tar_file_key}.tar'
3163
3643
 
3164
3644
  temp_dir = tempfile.mkdtemp()
@@ -3174,9 +3654,9 @@ class Ci(ExitStacked):
3174
3654
  os.makedirs(temp_requirements_dir)
3175
3655
 
3176
3656
  download_requirements(
3177
- self.resolve_ci_image(),
3657
+ await self.resolve_ci_image(),
3178
3658
  temp_requirements_dir,
3179
- requirements_txts,
3659
+ self.requirements_txts(),
3180
3660
  )
3181
3661
 
3182
3662
  if self._file_cache is not None:
@@ -3193,16 +3673,16 @@ class Ci(ExitStacked):
3193
3673
 
3194
3674
  return temp_requirements_dir
3195
3675
 
3196
- @cached_nullary
3197
- def resolve_requirements_dir(self) -> str:
3676
+ @async_cached_nullary
3677
+ async def resolve_requirements_dir(self) -> str:
3198
3678
  with log_timing_context('Resolve requirements dir') as ltc:
3199
- requirements_dir = self._resolve_requirements_dir()
3679
+ requirements_dir = await self._resolve_requirements_dir()
3200
3680
  ltc.set_description(f'Resolve requirements dir: {requirements_dir}')
3201
3681
  return requirements_dir
3202
3682
 
3203
3683
  #
3204
3684
 
3205
- def _run_compose_(self) -> None:
3685
+ async def _run_compose_(self) -> None:
3206
3686
  setup_cmds = [
3207
3687
  'pip install --root-user-action ignore --find-links /requirements --no-index uv',
3208
3688
  (
@@ -3220,37 +3700,39 @@ class Ci(ExitStacked):
3220
3700
 
3221
3701
  #
3222
3702
 
3223
- with DockerComposeRun(DockerComposeRun.Config(
3703
+ async with DockerComposeRun(DockerComposeRun.Config(
3224
3704
  compose_file=self._cfg.compose_file,
3225
3705
  service=self._cfg.service,
3226
3706
 
3227
- image=self.resolve_ci_image(),
3707
+ image=await self.resolve_ci_image(),
3228
3708
 
3229
3709
  cmd=ci_cmd,
3230
3710
 
3231
3711
  run_options=[
3232
3712
  '-v', f'{os.path.abspath(self._cfg.project_dir)}:/project',
3233
- '-v', f'{os.path.abspath(self.resolve_requirements_dir())}:/requirements',
3713
+ '-v', f'{os.path.abspath(await self.resolve_requirements_dir())}:/requirements',
3234
3714
  ],
3235
3715
 
3236
3716
  cwd=self._cfg.project_dir,
3717
+
3718
+ no_dependencies=self._cfg.no_dependencies,
3237
3719
  )) as ci_compose_run:
3238
- ci_compose_run.run()
3720
+ await ci_compose_run.run()
3239
3721
 
3240
- def _run_compose(self) -> None:
3722
+ async def _run_compose(self) -> None:
3241
3723
  with log_timing_context('Run compose'):
3242
- self._run_compose_()
3724
+ await self._run_compose_()
3243
3725
 
3244
3726
  #
3245
3727
 
3246
- def run(self) -> None:
3247
- self.load_compose_service_dependencies()
3728
+ async def run(self) -> None:
3729
+ await self.load_compose_service_dependencies()
3248
3730
 
3249
- self.resolve_ci_image()
3731
+ await self.resolve_ci_image()
3250
3732
 
3251
- self.resolve_requirements_dir()
3733
+ await self.resolve_requirements_dir()
3252
3734
 
3253
- self._run_compose()
3735
+ await self._run_compose()
3254
3736
 
3255
3737
 
3256
3738
  ########################################
@@ -3266,7 +3748,7 @@ class GithubCli(ArgparseCli):
3266
3748
  argparse_arg('key'),
3267
3749
  )
3268
3750
  def get_cache_entry(self) -> None:
3269
- shell_client = GithubV1CacheShellClient()
3751
+ shell_client = GithubCacheServiceV1ShellClient()
3270
3752
  entry = shell_client.run_get_entry(self.args.key)
3271
3753
  if entry is None:
3272
3754
  return
@@ -3325,18 +3807,21 @@ class CiCli(ArgparseCli):
3325
3807
  argparse_arg('--docker-file'),
3326
3808
  argparse_arg('--compose-file'),
3327
3809
  argparse_arg('-r', '--requirements-txt', action='append'),
3810
+
3328
3811
  argparse_arg('--github-cache', action='store_true'),
3329
3812
  argparse_arg('--cache-dir'),
3813
+
3330
3814
  argparse_arg('--always-pull', action='store_true'),
3815
+ argparse_arg('--always-build', action='store_true'),
3816
+
3817
+ argparse_arg('--no-dependencies', action='store_true'),
3331
3818
  )
3332
3819
  async def run(self) -> None:
3333
3820
  project_dir = self.args.project_dir
3334
3821
  docker_file = self.args.docker_file
3335
3822
  compose_file = self.args.compose_file
3336
- service = self.args.service
3337
3823
  requirements_txts = self.args.requirements_txt
3338
3824
  cache_dir = self.args.cache_dir
3339
- always_pull = self.args.always_pull
3340
3825
 
3341
3826
  #
3342
3827
 
@@ -3406,14 +3891,14 @@ class CiCli(ArgparseCli):
3406
3891
 
3407
3892
  #
3408
3893
 
3409
- with Ci(
3894
+ async with Ci(
3410
3895
  Ci.Config(
3411
3896
  project_dir=project_dir,
3412
3897
 
3413
3898
  docker_file=docker_file,
3414
3899
 
3415
3900
  compose_file=compose_file,
3416
- service=service,
3901
+ service=self.args.service,
3417
3902
 
3418
3903
  requirements_txts=requirements_txts,
3419
3904
 
@@ -3422,12 +3907,15 @@ class CiCli(ArgparseCli):
3422
3907
  'python3 -m pytest -svv test.py',
3423
3908
  ])),
3424
3909
 
3425
- always_pull=always_pull,
3910
+ always_pull=self.args.always_pull,
3911
+ always_build=self.args.always_build,
3912
+
3913
+ no_dependencies=self.args.no_dependencies,
3426
3914
  ),
3427
3915
  file_cache=file_cache,
3428
3916
  shell_cache=shell_cache,
3429
3917
  ) as ci:
3430
- ci.run()
3918
+ await ci.run()
3431
3919
 
3432
3920
 
3433
3921
  async def _async_main() -> ta.Optional[int]: